package producer

import (
	"encoding/json"
	"fmt"
	"math"
	"os"
	"strings"

	"github.com/duke-git/lancet/v2/mathutil"
	"github.com/paulmach/orb"
	"github.com/paulmach/orb/geojson"

	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: 9,
		gridsize:       19,
		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

	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,
	)

	if rpc.minH < -10.0 {
		rpc.minH = 0.0
	}

	if rpc.maxH < -10.0 {
		rpc.maxH = 0.0
	}

	rpc.heightOffset = (rpc.minH + rpc.maxH) / 2.0
}

// 虚拟控制点
func (rpc *RPC) generateVirtualGCP() {
	log.Infof("Generating virtual GCPs, %d x %d x %d",
		rpc.gridsize+1, rpc.gridsize+1, rpc.elevationLayer+1)
	points := gridImage2(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)
	rpc.saveGCP()

	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)
	}

	rpc.saveVec(strings.Replace(rpc.scene.Tiff, ".tiff", ".vec.txt", -1),
		rowVec, colVec, latVec, lonVec, heightVec)

	rpc.lineOffset = float64(rpc.scene.Height / 2)
	rpc.lineScale = float64(rpc.scene.Height)
	rpc.sampOffset = float64(rpc.scene.Width / 2)
	rpc.sampScale = float64(rpc.scene.Width)
	rowVec = normalize2(rowVec, rpc.lineOffset, rpc.lineScale)
	colVec = normalize2(colVec, rpc.sampOffset, rpc.sampScale)
	// rowVec.ScaleVec(1.0/rpc.lineScale, rowVec)
	// colVec.ScaleVec(1.0/rpc.sampScale, colVec)

	// rowVec, rpc.lineOffset, rpc.lineScale = normalize(rowVec)
	// colVec, rpc.sampOffset, rpc.sampScale = normalize(colVec)
	latVec, rpc.latOffset, rpc.latScale = normalize(latVec)
	lonVec, rpc.longOffset, rpc.longScale = normalize(lonVec)
	heightVec, rpc.heightOffset, rpc.heightScale = normalize(heightVec)

	// rpc.latOffset, rpc.latScale = (rpc.maxLat+rpc.minLat)/2.0, (rpc.maxLat - rpc.minLat)
	// rpc.longOffset, rpc.longScale = (rpc.maxLon+rpc.minLon)/2.0, (rpc.maxLon - rpc.minLon)
	// rpc.heightOffset, rpc.heightScale = (rpc.maxH+rpc.minH)/2.0+500.0, 500.0

	rpc.saveVec(strings.Replace(rpc.scene.Tiff, ".tiff", ".vec_norm.txt", -1),
		rowVec, colVec, latVec, lonVec, heightVec)

	// 设计矩阵 B = [ 20个分子系数 19个分母系数 ]
	B := buildDesignMatrix(rowVec, 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]
	}

	// 列参数
	D := buildDesignMatrix(colVec, latVec, lonVec, heightVec)
	K, err := SolveNormalEquation(D, 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 normalize2(v *mat.VecDense, vOff, vScale float64) *mat.VecDense {
	for i := 0; i < v.Len(); i++ {
		v.SetVec(i, (v.AtVec(i)-vOff)/vScale)
	}

	return v
}

func buildDesignMatrix(vec, 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)
		r_c := vec.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*r_c)
		B.Set(i, 21, -P*r_c)
		B.Set(i, 22, -H*r_c)
		B.Set(i, 23, -L*P*r_c)
		B.Set(i, 24, -L*H*r_c)
		B.Set(i, 25, -P*H*r_c)
		B.Set(i, 26, -L*L*r_c)
		B.Set(i, 27, -P*P*r_c)
		B.Set(i, 28, -H*H*r_c)
		B.Set(i, 29, -P*L*H*r_c)
		B.Set(i, 30, -L*L*L*r_c)
		B.Set(i, 31, -L*P*P*r_c)
		B.Set(i, 32, -L*H*H*r_c)
		B.Set(i, 33, -L*L*P*r_c)
		B.Set(i, 34, -P*P*P*r_c)
		B.Set(i, 35, -P*H*H*r_c)
		B.Set(i, 36, -L*L*H*r_c)
		B.Set(i, 37, -P*P*H*r_c)
		B.Set(i, 38, -H*H*H*r_c)
	}

	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 {
		// 岭估计方法调整法方程状态，使得矩阵非奇异，最小二乘平差可以收敛
		r, c := At.Dims()
		log.Infof("cannot inverse design matrix(%d*%d): %v", r, c, err)
		log.Info("try to adjust design matrix with +kI, k=0.0000001")
		k := 0.0000001 // [0.00000005, 0.000005]
		I := mat.NewDiagDense(r, nil)
		for i := 0; i < r; i++ {
			I.SetDiag(i, k)
		}
		At.Add(&At, I)

		err = AtInv.Inverse(&At)
	}

	if err != nil {
		log.Infof("cannot inverse design matrix: %v, try SVD method", err)
		// 计算矩阵的 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)
	if err != nil {
		log.Errorf("Failed to write RPC model file: %v", err)
	}
	f.Sync()
	return err
}

func (rpc *RPC) saveGCP() error {
	name := strings.Replace(rpc.scene.Tiff, ".tiff", ".gcp.geojson", -1)
	log.Infof("save gcp to %s", name)
	f, err := os.Create(name)
	if err != nil {
		log.Errorf("Failed to create GCP file: %v", err)
		return err
	}
	defer f.Close()

	var gcp geojson.FeatureCollection
	for _, p := range rpc.GCPs {
		point := orb.Point{p.L, p.P}
		feature := geojson.NewFeature(point)
		feature.Properties = map[string]interface{}{
			"H": p.H,
			"Y": p.Y,
			"X": p.X,
		}
		gcp.Features = append(gcp.Features, feature)
	}
	data, _ := json.Marshal(gcp)
	f.Write(data)
	f.Sync()

	return nil
}

func (rpc *RPC) saveVec(name string, rowVec, colVec, latVec, lonVec, heightVec *mat.VecDense) error {
	f, err := os.Create(name)

	if err != nil {
		log.Errorf("Failed to create vec file: %v", err)
		return err
	}
	defer f.Close()

	for i := 0; i < rowVec.Len(); i++ {
		f.WriteString(fmt.Sprintf("%.8f\t%.8f\t%.8f\t%.8f\t%.8f\n",
			rowVec.AtVec(i), colVec.AtVec(i), latVec.AtVec(i), lonVec.AtVec(i), heightVec.AtVec(i)))
	}
	return nil
}