package main

import (
	"os"
	"path/filepath"
	"strings"

	"github.com/airbusgeo/godal"
	"github.com/spf13/cobra"
	imageproc "starwiz.cn/sjy01/image-proc"
)

var (
	PanRawFile           string
	MssRawFile           string
	AuxRawFile           string
	SaveRegisteredMssRaw bool
	PansharpenIHS        bool
	OutputDir            string
)

var procCmd = &cobra.Command{
	Use:   "proc",
	Short: "process images",
	Long:  `process images`,
	Run: func(cmd *cobra.Command, args []string) {
		reg := imageproc.NewRegistrator(imageproc.DownSampled)
		reg.Params = imageproc.Params{
			PanRawFile:           PanRawFile,
			MssRawFile:           MssRawFile,
			AuxRawFile:           AuxRawFile,
			SaveRegisteredMssRaw: SaveRegisteredMssRaw,
			PansharpenIHS:        PansharpenIHS,
			OutputDir:            OutputDir,
			MssTiffFile: filepath.Join(OutputDir,
				strings.TrimSuffix(filepath.Base(MssRawFile),
					filepath.Ext(MssRawFile))+".tiff"),
			PanTiffFile: filepath.Join(OutputDir,
				strings.TrimSuffix(filepath.Base(PanRawFile),
					filepath.Ext(PanRawFile))+".tiff"),
		}
		reg.Params.FusTIffFile = strings.Replace(reg.Params.MssTiffFile, ".tiff", "_FUS.tiff", 1)

		if err := reg.LoadMssRaw(); err != nil {
			panic(err)
		}

		if err := reg.LoadPanRaw(); err != nil {
			panic(err)
		}

		godal.RegisterAll()
		os.MkdirAll(OutputDir, 0755)

		if err := reg.DoPhaseCorrelation(); err != nil {
			panic(err)
		}
		reg.DoCoRegestration()

		if SaveRegisteredMssRaw {
			registerdMSSRAW := filepath.Join(
				OutputDir,
				strings.TrimSuffix(filepath.Base(MssRawFile), filepath.Ext(MssRawFile))+"_registered.RAW",
			)
			reg.SaveRegisteredMssToRaw(registerdMSSRAW)
		}

		reg.SaveOriginalPanToGDALGTiff(reg.Params.PanTiffFile)
		reg.SaveRegisteredMssToGDALGTiff(reg.Params.MssTiffFile)

		if PansharpenIHS {
			reg.SavePansharpenedToGDALGTiff(reg.Params.FusTIffFile)
		}

		reg.Clean()
	},
}

func init() {
	rootCmd.AddCommand(procCmd)

	procCmd.Flags().StringVarP(&PanRawFile, "pan", "p", "pan.raw", "PAN image raw file path")
	procCmd.Flags().StringVarP(&MssRawFile, "mss", "m", "mss.raw", "MSS image raw file path")
	procCmd.Flags().StringVarP(&AuxRawFile, "aux", "a", "pms.aux", "AUX image raw file path")
	procCmd.Flags().BoolVarP(&SaveRegisteredMssRaw, "srmss", "s", false, "save registered MSS image raw file")
	procCmd.Flags().BoolVarP(&PansharpenIHS, "fus", "", false, "pansharpen using IHS")
	procCmd.Flags().StringVarP(&OutputDir, "output-dir", "o", "data", "output directory")
}