package producer

import (
	"fmt"
	"math"

	log "github.com/sirupsen/logrus"
	"gonum.org/v1/gonum/mat"
)

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 solveCoefficients(f, latVec, lonVec, heightVec *mat.VecDense) ([]float64, error) {
	M := setupSystemOfEquations(f, latVec, lonVec, heightVec)
	n := f.Len()
	weights := mat.NewDiagDense(n, nil)
	for i := 0; i < n; i++ {
		weights.SetDiag(i, 1.0)
	}
	w2 := mat.NewDiagDense(n, nil)

	iterations := 0
	var x mat.VecDense
	// var e0 float64

	for iterations < 20 {
		iterations++

		// w2 = weights^2
		for i := 0; i < n; i++ {
			w2.SetDiag(i, weights.At(i, i)*weights.At(i, i))
		}

		// x = (M^T * w2 * M)^-1 * M^T * w2 * R
		var MtW2 mat.Dense
		MtW2.Mul(M.T(), w2)
		var MtW2M mat.Dense
		MtW2M.Mul(&MtW2, M)

		invMtW2M, err := invertRPCMatrix(&MtW2M)
		if err != nil {
			return nil, err
		}
		var MtW2F mat.VecDense
		MtW2F.MulVec(&MtW2, f)

		x.MulVec(invMtW2M, &MtW2F)

		numerator := mat.NewVecDense(20, nil)
		denominator := mat.NewVecDense(20, nil)
		numerator.SetVec(0, x.AtVec(0))
		denominator.SetVec(0, 1.0)
		for idx := 1; idx < 20; idx++ {
			numerator.SetVec(idx, x.AtVec(idx))
			denominator.SetVec(idx, x.AtVec(idx+19))
		}

		weights = setupWeightMatrix(denominator, latVec, lonVec, heightVec)

		// residual = m.t()*w2*(m*tempCoeff-r);
		// var temp0, temp1, residual mat.VecDense
		// temp0.MulVec(M, &x)
		// temp1.SubVec(&temp0, f)
		// residual.MulVec(&MtW2, &temp1)
		// var square mat.Dense
		// square.Mul(residual.T(), &residual)
		// residualValue := math.Sqrt(square.At(0, 0))
		// if residualValue < 0.0001 {
		// 	break
		// }

		// fmt.Printf("residual value: %.16f\n", residualValue)
		// fmt.Printf("iterations: %d\n", iterations)

		e := 0.0
		for i := 0; i < n; i++ {
			Rcal := project(numerator, denominator, latVec.AtVec(i), lonVec.AtVec(i), heightVec.AtVec(i))
			e += math.Pow(Rcal-f.AtVec(i), 2)
		}
		e = math.Sqrt(e / float64(n))
		fmt.Println("iterations:", iterations, "r error:", e)

		if e < 1e-5 {
			break
		}

		// dnum := mat.NewVecDense(20, nil)
		// b0, b1 := true, true
		// dden := mat.NewVecDense(19, nil)
		// for i := 0; i < n; i++ {
		// 	dnum.SetVec(i, math.Abs(numerator.AtVec(i)-numerator0.AtVec(i)))
		// 	numerator0.SetVec(i, numerator.AtVec(i))
		// 	fmt.Println("dnum:", i, dnum.AtVec(i))
		// 	if dnum.AtVec(i) > 0.000001 {
		// 		b0 = false
		// 		break
		// 	}
		// }
		// for i := 1; i < n; i++ {
		// 	dden.SetVec(i, math.Abs(denominator.AtVec(i)-denominator0.AtVec(i)))
		// 	denominator0.SetVec(i, denominator.AtVec(i))
		// 	if dden.AtVec(i) > 0.000001 {
		// 		b1 = false
		// 		break
		// 	}
		// }

		// if b0 && b1 {
		// 	break
		// }
	}

	log.Println("iterations:", iterations)

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

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

	return B
}

// 构建权矩阵 [ 1/B ]
func setupWeightMatrix(coeffs, latVec, lonVec, heightVec *mat.VecDense) *mat.DiagDense {
	n := latVec.Len()
	row := mat.NewDense(n, 20, nil)
	result := mat.NewDiagDense(n, nil)
	for i := 0; i < n; i++ {
		P := latVec.AtVec(i)
		L := lonVec.AtVec(i)
		H := heightVec.AtVec(i)

		row.Set(i, 0, 1)
		row.Set(i, 1, L)
		row.Set(i, 2, P)
		row.Set(i, 3, H)
		row.Set(i, 4, L*P)
		row.Set(i, 5, L*H)
		row.Set(i, 6, P*H)
		row.Set(i, 7, L*L)
		row.Set(i, 8, P*P)
		row.Set(i, 9, H*H)
		row.Set(i, 10, P*L*H)
		row.Set(i, 11, L*L*L)
		row.Set(i, 12, L*P*P)
		row.Set(i, 13, L*H*H)
		row.Set(i, 14, L*L*P)
		row.Set(i, 15, P*P*P)
		row.Set(i, 16, P*H*H)
		row.Set(i, 17, L*L*H)
		row.Set(i, 18, P*P*H)
		row.Set(i, 19, H*H*H)

		var B float64
		for idx2 := 0; idx2 < 20; idx2++ {
			B += coeffs.AtVec(idx2) * row.At(i, idx2)
		}

		result.SetDiag(i, 1/B)
	}

	return result
}

func invertRPCMatrix(At *mat.Dense) (*mat.Dense, error) {
	var AtInv mat.Dense
	err := AtInv.Inverse(At)

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

	return &AtInv, 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 {
		// 岭估计方法调整法方程状态，使得矩阵非奇异，最小二乘平差可以收敛
		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 localize(num, den *mat.VecDense, row, col float64) (P, L, H float64) {

	return
}

func project(num, den *mat.VecDense, P, L, H float64) (v float64) {
	v = applyPoly(num, P, L, H) / applyPoly(den, P, L, H)
	return v
}

func applyPoly(poly *mat.VecDense,
	P, L, H float64) (v float64) {
	v = 0.0
	v += poly.AtVec(0)
	v += poly.AtVec(1) * L
	v += poly.AtVec(2) * P
	v += poly.AtVec(3) * H
	v += poly.AtVec(4) * L * P
	v += poly.AtVec(5) * L * H
	v += poly.AtVec(6) * P * H
	v += poly.AtVec(7) * L * L
	v += poly.AtVec(8) * P * P
	v += poly.AtVec(9) * H * H
	v += poly.AtVec(10) * P * L * H
	v += poly.AtVec(11) * L * L * L
	v += poly.AtVec(12) * L * P * P
	v += poly.AtVec(13) * L * H * H
	v += poly.AtVec(14) * L * L * P
	v += poly.AtVec(15) * P * P * P
	v += poly.AtVec(16) * P * H * H
	v += poly.AtVec(17) * L * L * H
	v += poly.AtVec(18) * P * P * H
	v += poly.AtVec(19) * H * H * H
	return v
}