package rrc

import (
	"bufio"
	"fmt"
	"os"
	"path/filepath"
	"runtime"
	"sync"

	log "github.com/sirupsen/logrus"
	"gocv.io/x/gocv"
	"starwiz.cn/sjy01/image-proc/pkg/payload"
	"starwiz.cn/sjy01/image-proc/pkg/utils"
)

// Relative Radiation Correction
// 采用在轨统计定标方法。利用卫星在轨后获取的常规影像数据，统计每个探元出现的灰度频次。

const (
	PANCameraProbeNum = payload.PAN_PIXEL_WIDTH // 全色探元数
	MSSCameraProbeNum = payload.MSS_PIXEL_WIDTH // 多光谱探元数
	MaxGrayLevel      = 65536                   // 16bit像素值
)

type RRC struct {
	LUTOutDir  string
	Histograms [5]ProbeHistogram
}

func NewRRC(dir string) *RRC {
	r := RRC{LUTOutDir: dir}
	r.Histograms[0].init(PANCameraProbeNum)
	for i := 1; i < 5; i++ {
		r.Histograms[i].init(MSSCameraProbeNum)
	}

	os.MkdirAll(r.LUTOutDir, 0755)

	return &r
}

func (rrc *RRC) Statistical(dsPAN, dsMSS string) error {
	rrc.StatisticalPAN(dsPAN)
	rrc.StatisticalMSS(dsMSS)

	return nil
}

func (rrc *RRC) Close() {}

// 基于探元直方图匹配的CBERS01星CCD数据相对辐射校正方法 [潘志强, 2005]
// 统计探元灰度的累积概率密度
func (rrc *RRC) StatisticalPAN(dsfile string) error {
	utils.PrintMemStats()
	f, err := os.Open(dsfile)
	if err != nil {
		return err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	var files []string
	for scanner.Scan() {
		files = append(files, scanner.Text())
	}

	// 并发处理
	var wg sync.WaitGroup
	jobs := make(chan struct{}, 5)

	var mutex sync.Mutex
	for _, file := range files {
		wg.Add(1)
		go func(l0 string) {
			defer wg.Done()

			jobs <- struct{}{}
			defer func() { <-jobs }()

			log.Println("statistical PAN RAW: ", l0)
			data, err := os.ReadFile(l0)
			if err != nil {
				log.Error("Error reading file: ", err)
				return
			}

			height := len(data) / (PANCameraProbeNum * 2)
			img, err := gocv.NewMatFromBytes(height, PANCameraProbeNum, gocv.MatTypeCV16U, data)
			if err != nil {
				log.Error("Error creating Mat from bytes: ", err)
				return
			}

			var hist ProbeHistogram
			hist.init0(PANCameraProbeNum)
			hist.statistical(img)
			img.Close()

			mutex.Lock()

			log.Println("sum PAN histogram...")
			rrc.Histograms[0].sum([]*ProbeHistogram{&hist})
			hist.free()
			runtime.GC()
			utils.PrintMemStats()

			mutex.Unlock()

		}(file)
	}

	wg.Wait()

	log.Println("compute PAN histogram...")
	rrc.Histograms[0].compute()
	log.Println("save PAN gray table matrix.")
	rrc.Histograms[0].saveLUT(filepath.Join(rrc.LUTOutDir, "B0.LUT"))

	return nil
}

func (rrc *RRC) StatisticalMSS(dsfile string) error {
	utils.PrintMemStats()
	f, err := os.Open(dsfile)
	if err != nil {
		return err
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)

	var files []string
	for scanner.Scan() {
		files = append(files, scanner.Text())
	}

	var wg sync.WaitGroup
	jobs := make(chan struct{}, 5)

	var mutex sync.Mutex
	for _, file := range files {
		wg.Add(1)
		go func(l0 string) {
			defer wg.Done()
			jobs <- struct{}{}
			defer func() { <-jobs }()

			log.Println("statistical MSS RAW: ", l0)
			data, err := os.ReadFile(l0)
			if err != nil {
				log.Error("Error reading file: ", err)
				return
			}

			width := MSSCameraProbeNum
			height := len(data) / (width * 4 * 2)
			mssData := make([][]byte, 4)
			for h := 0; h < height; h++ {
				row := data[h*width*4*2 : (h+1)*width*4*2]
				for i := 0; i < 4; i++ {
					mssData[i] = append(mssData[i], row[i*width*2:(i+1)*width*2]...)
				}
			}

			var mssImages [4]gocv.Mat
			for i := 0; i < 4; i++ {
				mssImages[i], err = gocv.NewMatFromBytes(height, width, gocv.MatTypeCV16U, mssData[i])
				if err != nil {
					log.Error("Error creating Mat from bytes: ", err)
					return
				}

				var hist ProbeHistogram
				hist.init0(MSSCameraProbeNum)
				hist.statistical(mssImages[i])
				mssImages[i].Close()

				mutex.Lock()

				log.Println("sum MSS histogram...")
				rrc.Histograms[i+1].sum([]*ProbeHistogram{&hist})
				hist.free()
				runtime.GC()
				utils.PrintMemStats()

				mutex.Unlock()
			}
		}(file)
	}

	wg.Wait()
	for i := 1; i < 5; i++ {
		log.Println("compute MSS histogram...")
		rrc.Histograms[i].compute()
		log.Println("save MSS gray table matrix.")
		rrc.Histograms[i].saveLUT(filepath.Join(rrc.LUTOutDir, fmt.Sprintf("B%d.LUT", i)))
	}

	return nil
}