package rrc

import (
	"fmt"
	"math"
	"os"

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

type ProbeHistogram struct {
	probes int         // 探元数
	N_i    []int64     // N_i 探元像素总数 PAN 0-9343 MSS 0-2335
	n_ik   [][]int64   // n_ik 第i探元灰度等级为k的像素数统计 PAN 9343 x 65536  MSS 2335 x 65536
	p_ik   [][]float64 // p_ik = n_ik / N_i 探元灰度概率密度 PAN 9343 x 65536  MSS 2335 x 65536
	m_l    []int64     // 具有灰度等级l的像素总数 l = 0-65535
	M      int64       // 参与直方图统计的总像素数
	P_l    []float64   // P_l = m_l/M 所有探元的期望直方图灰度等级为l的概率密度
	S_ik   [][]float64 // S_ik = sum(p_ij),j=0..k 第i个探元直方图灰度等级k的累积概率密度
	V_l    []float64   // V_l = sum(P_j),j=0..l // 期望直方图灰度级l对应的累积概率密度
	Tmat   *mat.Dense  // 第i个像元的j灰度等级对应的新的灰度值，用于修正图像
}

func (hist *ProbeHistogram) init(width int) {
	hist.probes = width
	hist.M = 0
	hist.N_i = make([]int64, width)
	hist.n_ik = make([][]int64, width)
	hist.p_ik = make([][]float64, width)
	hist.m_l = make([]int64, MaxGrayLevel)
	hist.P_l = make([]float64, MaxGrayLevel)
	hist.S_ik = make([][]float64, width)
	hist.V_l = make([]float64, MaxGrayLevel)

	for i := 0; i < width; i++ {
		hist.n_ik[i] = make([]int64, MaxGrayLevel)
		hist.p_ik[i] = make([]float64, MaxGrayLevel)
		hist.S_ik[i] = make([]float64, MaxGrayLevel)
	}

	hist.Tmat = mat.NewDense(width, MaxGrayLevel, nil)
}

func (hist *ProbeHistogram) statistical(img gocv.Mat) error {
	hist.M += int64(img.Rows() * img.Cols())
	fmt.Println("Hist.M:", hist.M)

	// 探元i像素总数
	for i := 0; i < hist.probes; i++ {
		hist.N_i[i] += int64(img.Rows())
	}

	// 探元i灰度等级k的像素数统计
	for y := 0; y < img.Rows(); y++ {
		for x := 0; x < img.Cols(); x++ {
			gray := uint16(img.GetShortAt(x, y))
			hist.n_ik[x][gray]++
		}
	}

	// 灰度等级l的像素总数
	for gray := 0; gray < MaxGrayLevel; gray++ {
		for i := 0; i < hist.probes; i++ {
			hist.m_l[gray] += int64(hist.n_ik[i][gray])
		}
	}

	return nil
}

func (hist *ProbeHistogram) compute() {
	// 探元i灰度概率密度
	for i := 0; i < hist.probes; i++ {
		for k := 0; k < MaxGrayLevel; k++ {
			hist.p_ik[i][k] = float64(hist.n_ik[i][k]) / float64(hist.N_i[i])
		}
	}

	// 所有探元的期望直方图灰度等级为l的概率密度
	for gray := 0; gray < MaxGrayLevel; gray++ {
		hist.P_l[gray] = float64(hist.m_l[gray]) / float64(hist.M)
	}

	// 第i个探元直方图灰度等级k的累积概率密度
	for i := 0; i < hist.probes; i++ {
		for k := 0; k < MaxGrayLevel; k++ {
			hist.S_ik[i][k] = 0
			for j := 0; j <= k; j++ {
				hist.S_ik[i][k] += hist.p_ik[i][j]
			}
		}
	}

	// 期望直方图灰度级l对应的累积概率密度
	for gray := 0; gray < MaxGrayLevel; gray++ {
		hist.V_l[gray] = 0
		for j := 0; j <= gray; j++ {
			hist.V_l[gray] += hist.P_l[j]
		}
	}

	// 生成查找表
	nT := 0
	for i := 0; i < hist.probes; i++ {
		for k := 0; k < MaxGrayLevel; k++ {
			for l := 0; l < MaxGrayLevel-1; l++ {
				if hist.S_ik[i][k] >= hist.V_l[l] && hist.S_ik[i][k] <= hist.V_l[l+1] {
					nT += 1
					if math.Abs(hist.S_ik[i][k]-hist.V_l[l]) <= math.Abs(hist.S_ik[i][k]-hist.V_l[l+1]) {
						hist.Tmat.Set(i, k, float64(l))
					} else {
						hist.Tmat.Set(i, k, float64(l+1))
					}
				}
			}
		}
	}

	if nT != hist.probes*MaxGrayLevel {
		logrus.Error("error in computing Tij table, some values are not satisfied")
	}
}

func (hist *ProbeHistogram) save(matrixFile string) error {
	log.Printf("total pixels: %d", hist.M)
	f, err := os.OpenFile(matrixFile, os.O_TRUNC|os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer f.Close()

	_, err = hist.Tmat.MarshalBinaryTo(f)
	if err != nil {
		return err
	}

	return nil
}

func LoadGrayTableMatrix(matrixFile string) (*mat.Dense, error) {
	f, err := os.OpenFile(matrixFile, os.O_RDONLY, 0644)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	matrix := mat.Dense{}
	if _, err := matrix.UnmarshalBinaryFrom(f); err != nil {
		return nil, err
	}

	return &matrix, nil
}

func (hist *ProbeHistogram) sum(hists []*ProbeHistogram) {
	for _, h := range hists {
		hist.M += h.M
		for i := 0; i < hist.probes; i++ {
			hist.N_i[i] += h.N_i[i]
			for k := 0; k < MaxGrayLevel; k++ {
				hist.n_ik[i][k] += h.n_ik[i][k]
			}
		}
		for gray := 0; gray < MaxGrayLevel; gray++ {
			hist.m_l[gray] += h.m_l[gray]
		}
	}
}