package producer

import (
	"fmt"
	"math"
	"os"
	"time"

	"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/auxilary"
	"starwiz.cn/sjy01/image-proc/pkg/calculator"
	"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()
	rpc.generateVirtualGCP()
	rpc.regularizeGCPs()

	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
	// for row := 0; row < rpc.scene.Height; row++ {
	// 	for col := 0; col < rpc.scene.Width; col++ {
	// 		p84 := rpc.calculateLatLonH(row, col)
	// 		rpc.GroundPoints = append(rpc.GroundPoints, &p84)

	// 		rpc.heightOffset += p84.H
	// 		rpc.latOffset += p84.P
	// 		rpc.longOffset += p84.L

	// 		if p84.H < rpc.minH {
	// 			rpc.minH = p84.H
	// 		}
	// 		if p84.H > rpc.maxH {
	// 			rpc.maxH = p84.H
	// 		}
	// 		if p84.P < rpc.minLat {
	// 			rpc.minLat = p84.P
	// 		}
	// 		if p84.P > rpc.maxLat {
	// 			rpc.maxLat = p84.P
	// 		}
	// 		if p84.L < rpc.minLon {
	// 			rpc.minLon = p84.L
	// 		}
	// 		if p84.L > rpc.maxLon {
	// 			rpc.maxLon = p84.L
	// 		}
	// 	}
	// }

	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

	rpc.calculateRegularizedParams()
}

// 虚拟控制点
func (rpc *RPC) generateVirtualGCP() {
	log.Info("Generating virtual GCPs...")
	// elevations := []float64{rpc.minH, (rpc.minH + rpc.maxH) / 2.0, rpc.maxH}
	deltaRow := float64(rpc.scene.Height) / float64(rpc.gridsize-1) // 像素平面Y方向步长
	deltaCol := float64(rpc.scene.Width) / float64(rpc.gridsize-1)  // 像素平面X方向步长

	for i := 0; i < rpc.gridsize; i++ {
		for j := 0; j < rpc.gridsize; j++ {
			imagePoint := ImagePoint{
				Y: deltaRow * float64(i),
				X: deltaCol * float64(j),
			}
			groudPoint84 := rpc.calculateLatLonH(int(imagePoint.Y), int(imagePoint.X))
			rpc.GCPs = append(rpc.GCPs, groudPoint84)
		}
	}
}

func (rpc *RPC) regularizeGCPs() {
	log.Info("Regularizing virtual GCPs...")
	for i := range rpc.GCPs {
		gcp := &rpc.GCPs[i]
		gcp.P = (gcp.P - rpc.latOffset) / rpc.latScale
		gcp.L = (gcp.L - rpc.longOffset) / rpc.longScale
		gcp.H = (gcp.H - rpc.heightOffset) / rpc.heightScale
		gcp.Y = (gcp.Y - rpc.lineOffset) / rpc.lineScale
		gcp.X = (gcp.X - rpc.sampOffset) / rpc.sampScale
	}
}

// 计算 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))
}

func (rpc *RPC) calculateLatLonH(row, col int) GroundPoint {
	auxIdx := rpc.registrator.sceneOffsetInAuxIndex(rpc.scene, row)
	as := rpc.registrator.AuxPlatforms[auxIdx]
	t := time.Unix(int64(auxilary.ReferenceTime2000)+int64(as.UTCTimeSec), int64(as.Microsecond)*1000).UTC()
	p84 := []float64{as.W84PosX, as.W84PosY, as.W84PosZ}

	Qsat2eci := calculator.Quaternion{W: as.QuatAttstarQ0, X: as.QuatAttstarQ1, Y: as.QuatAttstarQ2, Z: as.QuatAttstarQ3}
	groudPoint84, _ := calculator.IntersectionAttitude(Qsat2eci, p84, t, col)

	elv := dem.Dem1KmLT.Elevation(groudPoint84.Lon, groudPoint84.Lat)
	if elv < 0.0 {
		elv = 0.0
	}

	return GroundPoint{
		P: groudPoint84.Lat,
		L: groudPoint84.Lon,
		H: float64(elv),
		Y: float64(row),
		X: float64(col),
	}
}

// BuildDesignMatrix 构建设计矩阵
func (rpc *RPC) buildDesignMatrix(removeConstant bool) *mat.Dense {
	n := len(rpc.GCPs)
	numCols := 20
	if removeConstant {
		numCols = 19
	}
	A := mat.NewDense(n, numCols, nil)

	for i, pt := range rpc.GCPs {
		idx := 0
		if !removeConstant {
			A.Set(i, 0, 1)
		} else {
			idx = -1
		}

		A.Set(i, idx+1, pt.L)
		A.Set(i, idx+2, pt.P)
		A.Set(i, idx+3, pt.H)
		A.Set(i, idx+4, pt.L*pt.P)
		A.Set(i, idx+5, pt.L*pt.H)
		A.Set(i, idx+6, pt.P*pt.H)
		A.Set(i, idx+7, pt.L*pt.L)
		A.Set(i, idx+8, pt.P*pt.P)
		A.Set(i, idx+9, pt.H*pt.H)
		A.Set(i, idx+10, pt.P*pt.L*pt.H)
		A.Set(i, idx+11, pt.L*pt.L*pt.L)
		A.Set(i, idx+12, pt.L*pt.P*pt.P)
		A.Set(i, idx+13, pt.L*pt.H*pt.H)
		A.Set(i, idx+14, pt.L*pt.L*pt.P)
		A.Set(i, idx+15, pt.P*pt.P*pt.P)
		A.Set(i, idx+16, pt.P*pt.H*pt.H)
		A.Set(i, idx+17, pt.L*pt.L*pt.H)
		A.Set(i, idx+18, pt.P*pt.P*pt.H)
		A.Set(i, idx+19, pt.H*pt.H*pt.H)
	}
	return A
}

// SolveLeastSquares 使用最小二乘法求解RPC模型的系数
func (rpc *RPC) SolveLeastSquares() error {
	log.Info("Solving least squares...")
	n := len(rpc.GCPs)
	rowVec := mat.NewVecDense(n, nil)
	colVec := mat.NewVecDense(n, nil)

	for i, ip := range rpc.GCPs {
		rowVec.SetVec(i, ip.Y)
		colVec.SetVec(i, ip.X)
	}

	var rowNum, rowDen, colNum, colDen []float64
	var err error
	// 使用正规方程法求解分子系数
	A := rpc.buildDesignMatrix(false)
	log.Debug("solving normal equation for linenumcoef...")
	r, c := A.Dims()
	log.Debug("A:", r, c)
	// fmt.Printf("设计矩阵:\n%v\n", mat.Formatted(A, mat.Prefix(" "), mat.Excerpt(0)))
	rowNum, err = SolveNormalEquation(A, rowVec)
	if err != nil {
		return err
	}

	log.Debug("solving normal equation for samplenumcoef...")
	colNum, err = SolveNormalEquation(A, colVec)
	if err != nil {
		return err
	}

	// 对分母系数求解，添加固定项
	rowDen = make([]float64, 20)
	colDen = make([]float64, 20)

	rowDen[0] = 1.0
	colDen[0] = 1.0
	A_reduced := rpc.buildDesignMatrix(true)
	log.Debug("solving normal equation for linedemcoef...")
	rowDenCoeffs, err := SolveNormalEquation(A_reduced, rowVec)
	if err != nil {
		return err
	}
	copy(rowDen[1:], rowDenCoeffs)

	log.Debug("solving normal equation for sampledencoef...")
	colDenCoeffs, err := SolveNormalEquation(A_reduced, colVec)
	if err != nil {
		return err
	}
	copy(colDen[1:], colDenCoeffs)

	for i := 0; i < 20; i++ {
		rpc.LineCoef.NumCoefficients[i] = rowNum[i]
		rpc.LineCoef.DenCoefficients[i] = rowDen[i]
		rpc.SampleCoef.NumCoefficients[i] = colNum[i]
		rpc.SampleCoef.DenCoefficients[i] = colDen[i]
	}

	return nil
}

// 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 {
		return nil, fmt.Errorf("矩阵不可逆: %v", err)
	}

	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("%f,\n", rpc.LineCoef.NumCoefficients[i])
			lineDenCoef += fmt.Sprintf("%f,\n", rpc.LineCoef.DenCoefficients[i])
			sampNumCoef += fmt.Sprintf("%f,\n", rpc.SampleCoef.NumCoefficients[i])
			sampDenCoef += fmt.Sprintf("%f,\n", rpc.SampleCoef.DenCoefficients[i])
		} else {
			lineNumCoef += fmt.Sprintf("%f", rpc.LineCoef.NumCoefficients[i])
			lineDenCoef += fmt.Sprintf("%f", rpc.LineCoef.DenCoefficients[i])
			sampNumCoef += fmt.Sprintf("%f", rpc.SampleCoef.NumCoefficients[i])
			sampDenCoef += fmt.Sprintf("%f", rpc.SampleCoef.DenCoefficients[i])
		}
	}

	model := fmt.Sprintf(`satId = "SJY01";
bandId = "";
SpecId = "";
BEGIN_GROUP = IMAGE
	errBias = 1.0;
	errRand = 0.0;
	lineOffset = %f;
	sampOffset = %f;
	latOffset = %f;
	longOffset = %f;
	heightOffset = %f;
	lineScale = %f;
	sampScale = %f;
	latScale = %f;
	longScale = %f;
	heightScale = %f;
	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("Saving 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()
	fmt.Println(model)
	_, err = f.WriteString(model)
	return err
}