package producer

import (
	"fmt"
	"math"
	"os"

	"github.com/duke-git/lancet/v2/mathutil"
	"github.com/duke-git/lancet/v2/slice"

	log "github.com/sirupsen/logrus"
	"gonum.org/v1/gonum/mat"
	"starwiz.cn/sjy01/image-proc/pkg/dem"
)

type RPC struct {
	lineOffset, lineScale               float64
	sampOffset, sampScale               float64
	latOffset, longOffset, heightOffset float64
	latScale, longScale, heightScale    float64

	LineCoef   RPCModel
	SampleCoef RPCModel

	// GroundPoints                   []*GroundPoint
	minLat, maxLat, minLon, maxLon float64
	minH, maxH                     float64
	GCPs                           []GroundPoint
	elevationLayer                 int
	gridsize                       int

	scene       *Scene
	registrator *Registrator

	rpb string
}

// GroundPoint 表示地面点的三维坐标
type GroundPoint struct {
	P, L, H float64 // P-latitude, L-longitude, H-height
	Y, X    float64 // X-sample, Y-line
}

// ImagePoint 表示像素平面上的点
type ImagePoint struct {
	Y, X float64 // X-sample, Y-line
}

// RPCModel 包含20个系数的RPC模型
type RPCModel struct {
	NumCoefficients [20]float64 // 分子系数
	DenCoefficients [20]float64 // 分母系数
}

// rational polynomial coeffients
func NewRPC(r *Registrator, scene *Scene, rpb string) *RPC {
	rpc := RPC{
		elevationLayer: 3,
		gridsize:       20,
		registrator:    r,
		scene:          scene,
		rpb:            rpb,
	}

	log.Info("start RPC initialization for scene: ", scene.Tiff)
	rpc.init()

	return &rpc
}

// 初始化经纬度
func (rpc *RPC) init() {
	rpc.minH = 9999.0
	rpc.maxH = -9999.0
	rpc.minLat = 90.0
	rpc.maxLat = -90.0
	rpc.minLon = 180.0
	rpc.maxLon = -180.0

	rpc.minLat = mathutil.Min(rpc.scene.Meta.Corners.LowerLeft.Latitude,
		rpc.scene.Meta.Corners.LowerRight.Latitude,
		rpc.scene.Meta.Corners.UpperLeft.Latitude,
		rpc.scene.Meta.Corners.UpperRight.Latitude)
	rpc.maxLat = mathutil.Max(rpc.scene.Meta.Corners.LowerLeft.Latitude,
		rpc.scene.Meta.Corners.LowerRight.Latitude,
		rpc.scene.Meta.Corners.UpperLeft.Latitude,
		rpc.scene.Meta.Corners.UpperRight.Latitude)

	rpc.minLon = mathutil.Min(rpc.scene.Meta.Corners.LowerLeft.Longitude,
		rpc.scene.Meta.Corners.LowerRight.Longitude,
		rpc.scene.Meta.Corners.UpperLeft.Longitude,
		rpc.scene.Meta.Corners.UpperRight.Longitude)
	rpc.maxLon = mathutil.Max(rpc.scene.Meta.Corners.LowerLeft.Longitude,
		rpc.scene.Meta.Corners.LowerRight.Longitude,
		rpc.scene.Meta.Corners.UpperLeft.Longitude,
		rpc.scene.Meta.Corners.UpperRight.Longitude)

	rpc.latOffset = (rpc.minLat + rpc.maxLat) / 2.0
	rpc.longOffset = (rpc.minLon + rpc.maxLon) / 2.0

	var H []float64
	H = append(H, float64(dem.Dem1KmLT.Elevation(rpc.scene.Meta.Corners.LowerLeft.Longitude,
		rpc.scene.Meta.Corners.LowerLeft.Latitude)))
	H = append(H, float64(dem.Dem1KmLT.Elevation(rpc.scene.Meta.Corners.LowerRight.Longitude,
		rpc.scene.Meta.Corners.LowerRight.Latitude)))
	H = append(H, float64(dem.Dem1KmLT.Elevation(rpc.scene.Meta.Corners.UpperLeft.Longitude,
		rpc.scene.Meta.Corners.UpperLeft.Latitude)))
	H = append(H, float64(dem.Dem1KmLT.Elevation(rpc.scene.Meta.Corners.UpperRight.Longitude,
		rpc.scene.Meta.Corners.UpperRight.Latitude)))
	H = append(H, float64(dem.Dem1KmLT.Elevation(rpc.longOffset, rpc.latOffset)))
	slice.Sort(H, "asc")
	rpc.minH = H[0]
	rpc.maxH = H[len(H)-1]

	rpc.minH, rpc.maxH = dem.Dem1KmLT.MinMaxElevationInRect(
		rpc.scene.Meta.Corners.UpperLeft.Longitude,
		rpc.scene.Meta.Corners.UpperLeft.Latitude,
		rpc.scene.Meta.Corners.LowerRight.Longitude,
		rpc.scene.Meta.Corners.LowerRight.Latitude,
	)
	rpc.heightOffset = (rpc.minH + rpc.maxH) / 2.0
}

// 虚拟控制点
func (rpc *RPC) generateVirtualGCP() {
	log.Info("Generating virtual GCPs...")
	points := gridImage(rpc.gridsize, rpc.gridsize,
		rpc.scene.Height, rpc.scene.Width,
		rpc.elevationLayer, int(rpc.minH), int(rpc.maxH))

	for _, p := range points {
		p84 := rpc.registrator.calculateLatLonH(rpc.scene, p.Row, p.Col, p.H)
		rpc.GCPs = append(rpc.GCPs, GroundPoint{
			P: p84.Lat,
			L: p84.Lon,
			H: p84.H,
			Y: float64(p.Row),
			X: float64(p.Col),
		})
	}
}

func (rpc *RPC) RPC() error {
	rpc.generateVirtualGCP()
	n := len(rpc.GCPs)

	log.Info("num of virtual GCPs: ", n)

	rowVec := mat.NewVecDense(n, nil)
	colVec := mat.NewVecDense(n, nil)
	latVec := mat.NewVecDense(n, nil)
	lonVec := mat.NewVecDense(n, nil)
	heightVec := mat.NewVecDense(n, nil)

	for i, ip := range rpc.GCPs {
		rowVec.SetVec(i, ip.Y)
		colVec.SetVec(i, ip.X)
		latVec.SetVec(i, ip.P)
		lonVec.SetVec(i, ip.L)
		heightVec.SetVec(i, ip.H)
	}

	rowVec, rowOff, rowScale := normalize(rowVec)
	colVec, colOff, colScale := normalize(colVec)
	latVec, latOff, latScale := normalize(latVec)
	lonVec, lonOff, lonScale := normalize(lonVec)
	heightVec, heightOff, heightScale := normalize(heightVec)

	rpc.lineOffset = rowOff
	rpc.lineScale = rowScale
	rpc.sampOffset = colOff
	rpc.sampScale = colScale
	rpc.latOffset = latOff
	rpc.latScale = latScale
	rpc.longOffset = lonOff
	rpc.longScale = lonScale
	rpc.heightOffset = heightOff
	rpc.heightScale = heightScale

	// fmt.Printf("lineOffset: %f, lineScale: %f\n", rpc.lineOffset, rpc.lineScale)
	// fmt.Printf("sampOffset: %f, sampScale: %f\n", rpc.sampOffset, rpc.sampScale)
	// fmt.Printf("latOffset: %f, latScale: %f\n", rpc.latOffset, rpc.latScale)
	// fmt.Printf("longOffset: %f, longScale: %f\n", rpc.longOffset, rpc.longScale)
	// fmt.Printf("heightOffset: %f, heightScale: %f\n", rpc.heightOffset, rpc.heightScale)
	// fmt.Printf("X0: %f, Y0: %f\n", colVec.At(111, 0), rowVec.At(111, 0))
	// fmt.Printf("lat0: %f, lon0: %f, height0: %f\n", latVec.At(111, 0), lonVec.At(111, 0), heightVec.At(111, 0))

	// 设计矩阵 B = [ 20个分子系数 19个分母系数 ]
	B := buildDesignMatrix(latVec, lonVec, heightVec)

	// x = (B^T * B)^-1 * B^T * l, 其中 x = [a1..a20 b2..b20]^T
	// 行参数
	J, err := SolveNormalEquation(B, rowVec)
	if err != nil {
		return err
	}

	for i := 0; i < 20; i++ {
		rpc.LineCoef.NumCoefficients[i] = J[i]
	}
	rpc.LineCoef.DenCoefficients[0] = 1.0
	for i := 20; i < 39; i++ {
		rpc.LineCoef.DenCoefficients[i-19] = J[i]
	}

	// 列参数
	K, err := SolveNormalEquation(B, colVec)
	if err != nil {
		return err
	}

	for i := 0; i < 20; i++ {
		rpc.SampleCoef.NumCoefficients[i] = K[i]
	}
	rpc.SampleCoef.DenCoefficients[0] = 1.0
	for i := 20; i < 39; i++ {
		rpc.SampleCoef.DenCoefficients[i-19] = K[i]
	}
	return nil
}

func normalize(v *mat.VecDense) (*mat.VecDense, float64, float64) {
	var vOff, vScale float64
	vOff = mat.Sum(v) / float64(v.Len())
	vScale = math.Max(math.Abs(mat.Max(v)-vOff), math.Abs(mat.Min(v)-vOff))
	for i := 0; i < v.Len(); i++ {
		v.SetVec(i, (v.AtVec(i)-vOff)/vScale)
	}

	return v, vOff, vScale
}

func buildDesignMatrix(latVec, lonVec, heightVec *mat.VecDense) *mat.Dense {
	n := latVec.Len()
	// 设计矩阵 B = [ 20个分子系数 19个分母系数 ]
	B := mat.NewDense(n, 39, nil)
	for i := 0; i < n; i++ {
		P := latVec.AtVec(i)
		L := lonVec.AtVec(i)
		H := heightVec.AtVec(i)
		B.Set(i, 0, 1)
		B.Set(i, 1, L)
		B.Set(i, 2, P)
		B.Set(i, 3, H)
		B.Set(i, 4, L*P)
		B.Set(i, 5, L*H)
		B.Set(i, 6, P*H)
		B.Set(i, 7, L*L)
		B.Set(i, 8, P*P)
		B.Set(i, 9, H*H)
		B.Set(i, 10, P*L*H)
		B.Set(i, 11, L*L*L)
		B.Set(i, 12, L*P*P)
		B.Set(i, 13, L*H*H)
		B.Set(i, 14, L*L*P)
		B.Set(i, 15, P*P*P)
		B.Set(i, 16, P*H*H)
		B.Set(i, 17, L*L*H)
		B.Set(i, 18, P*P*H)
		B.Set(i, 19, H*H*H)
		B.Set(i, 20, -L)
		B.Set(i, 21, -P)
		B.Set(i, 22, -H)
		B.Set(i, 23, -L*P)
		B.Set(i, 24, -L*H)
		B.Set(i, 25, -P*H)
		B.Set(i, 26, -L*L)
		B.Set(i, 27, -P*P)
		B.Set(i, 28, -H*H)
		B.Set(i, 29, -P*L*H)
		B.Set(i, 30, -L*L*L)
		B.Set(i, 31, -L*P*P)
		B.Set(i, 32, -L*H*H)
		B.Set(i, 33, -L*L*P)
		B.Set(i, 34, -P*P*P)
		B.Set(i, 35, -P*H*H)
		B.Set(i, 36, -L*L*H)
		B.Set(i, 37, -P*P*H)
		B.Set(i, 38, -H*H*H)
	}

	return B
}

// 计算 RPC 正则化参数
func (rpc *RPC) calculateRegularizedParams() {
	rpc.lineOffset = float64(rpc.scene.Height) / 2.0
	rpc.sampOffset = float64(rpc.scene.Width) / 2.0
	rpc.lineScale = float64(rpc.scene.Height)
	rpc.sampScale = float64(rpc.scene.Width)

	// rpc.heightScale = math.Max(math.Abs(rpc.minH-rpc.heightOffset), math.Abs(rpc.maxH-rpc.heightOffset))
	rpc.heightScale = 500.0
	rpc.latScale = math.Max(math.Abs(rpc.minLat-rpc.latOffset), math.Abs(rpc.maxLat-rpc.latOffset))
	rpc.longScale = math.Max(math.Abs(rpc.minLon-rpc.longOffset), math.Abs(rpc.maxLon-rpc.longOffset))
}

// SolveNormalEquation 使用正规方程法求解最小二乘问题
func SolveNormalEquation(A *mat.Dense, b *mat.VecDense) ([]float64, error) {
	var At mat.Dense
	At.Mul(A.T(), A) // At = A^T * A

	// 求解 (A^T * A)^-1 * (A^T * b)
	var AtInv mat.Dense
	err := AtInv.Inverse(&At)
	if err != nil {
		// 计算矩阵的 SVD 分解
		var svd mat.SVD
		ok := svd.Factorize(&At, mat.SVDThin)
		if !ok {
			fmt.Println("SVD 分解失败")
			return nil, fmt.Errorf("设计矩阵不可逆, SVD 分解失败: %v", err)
		}

		// 获取 U、Σ 和 V^T
		var u, v mat.Dense
		svd.UTo(&u)
		svd.VTo(&v)
		sigma := svd.Values(nil)

		// 计算 Σ^+ (Sigma pseudo-inverse)
		m, n := u.Dims()
		sigmaInv := mat.NewDense(n, m, nil)
		for i := 0; i < len(sigma); i++ {
			if sigma[i] > 1e-10 { // 避免除以零
				sigmaInv.Set(i, i, 1/sigma[i])
			}
		}

		// 计算 V * Σ^+ * U^T
		var temp mat.Dense
		temp.Mul(&v, sigmaInv)
		AtInv.Mul(&temp, u.T())
	}

	var Atb mat.VecDense
	Atb.MulVec(A.T(), b) // Atb = A^T * b

	var x mat.VecDense
	x.MulVec(&AtInv, &Atb) // x = (A^T * A)^-1 * (A^T * b)

	return mat.Col(nil, 0, &x), nil
}

func (rpc *RPC) Output() string {
	var lineNumCoef, lineDenCoef, sampNumCoef, sampDenCoef string
	for i := 0; i < 20; i++ {
		if i < 19 {
			lineNumCoef += fmt.Sprintf("\t\t%.15e,\n", rpc.LineCoef.NumCoefficients[i])
			lineDenCoef += fmt.Sprintf("\t\t%.15e,\n", rpc.LineCoef.DenCoefficients[i])
			sampNumCoef += fmt.Sprintf("\t\t%.15e,\n", rpc.SampleCoef.NumCoefficients[i])
			sampDenCoef += fmt.Sprintf("\t\t%.15e,\n", rpc.SampleCoef.DenCoefficients[i])
		} else {
			lineNumCoef += fmt.Sprintf("\t\t%.15e", rpc.LineCoef.NumCoefficients[i])
			lineDenCoef += fmt.Sprintf("\t\t%.15e", rpc.LineCoef.DenCoefficients[i])
			sampNumCoef += fmt.Sprintf("\t\t%.15e", rpc.SampleCoef.NumCoefficients[i])
			sampDenCoef += fmt.Sprintf("\t\t%.15e", rpc.SampleCoef.DenCoefficients[i])
		}
	}

	model := fmt.Sprintf(`satId = "SJY01";
bandId = "";
SpecId = "";
BEGIN_GROUP = IMAGE
	errBias = 1.0;
	errRand = 0.0;
	lineOffset = %.8f;
	sampOffset = %.8f;
	latOffset = %.8f;
	longOffset = %.8f;
	heightOffset = %.8f;
	lineScale = %.8f;
	sampScale = %.8f;
	latScale = %.8f;
	longScale = %.8f;
	heightScale = %.8f;
	lineNumCoef = (
%s);
	lineDenCoef = (
%s);
	sampNumCoef = (
%s);
	sampDenCoef = (
%s);
END_GROUP = IMAGE
END;
`,
		rpc.lineOffset, rpc.sampOffset, rpc.latOffset, rpc.longOffset, rpc.heightOffset,
		rpc.lineScale, rpc.sampScale, rpc.latScale, rpc.longScale, rpc.heightScale,
		lineNumCoef, lineDenCoef, sampNumCoef, sampDenCoef,
	)

	return model
}

func (rpc *RPC) SaveRpb() error {
	log.Infof("save RPC model to %s", rpc.rpb)
	model := rpc.Output()
	f, err := os.Create(rpc.rpb)
	if err != nil {
		log.Errorf("Failed to create RPC model file: %v", err)
		return err
	}
	defer f.Close()
	_, err = f.WriteString(model)
	return err
}