'''
PANCO suite
Adriano Tullo - v.0.1 - 06/06/2024
INAF OAPd
'''

import numpy as np
import scipy.ndimage as sn
from osgeo import gdal, gdal_array
from numba import njit

def CaSSIS_ColorSimulator_Main(cas_input, band_NIR = 1, band_RED = 2, band_PAN = 3, band_BLU=4):

    outputName = f'{cas_input.split(".")[0]}_HCS.tif'

    #Computing weights

    cas_NIR = gdal_array.LoadFile(cas_input, band_list=[band_NIR])
    cas_NIR[cas_NIR > 5] = np.nan
    cas_NIR[np.isnan(cas_NIR)] = 0
    cas_NIR[cas_NIR < -1] = 0
    cas_RED = gdal_array.LoadFile(cas_input, band_list=[band_RED])
    cas_RED[cas_RED > 3] = np.nan
    cas_RED[np.isnan(cas_RED)] = 0
    cas_RED[cas_RED < -1] = 0
    cas_PAN = gdal_array.LoadFile(cas_input, band_list=[band_PAN])
    cas_PAN[cas_PAN > 3] = np.nan
    cas_PAN[np.isnan(cas_PAN)] = 0
    cas_PAN[cas_PAN < -1] = 0
    cas_BLU = gdal_array.LoadFile(cas_input, band_list=[band_BLU])
    cas_BLU[cas_BLU > 3] = np.nan
    cas_BLU[np.isnan(cas_BLU)] = 0
    cas_BLU[cas_BLU < -1] = 0

    #CORRETTO MEDIA INVECE DI SOMMA OPPURE VALUTARE MEDIA PESATA
    cas_IR = (cas_NIR + cas_RED)/2

    stacked_array = np.stack([cas_IR, cas_PAN, cas_BLU], axis=0)

    driver = gdal.GetDriverByName('GTiff')
    n, rows, cols = stacked_array.shape
    dataset = driver.Create(outputName, cols, rows, n,
                            gdal.GDT_Float32)

    for b in range(1, n + 1):
        band = dataset.GetRasterBand(b)  # GetRasterBand is not zero indexed
        band.WriteArray(stacked_array[b - 1])  # Numpy is zero indexed

    dataset = None
    stacked_array = None

    temp_cas_Open = gdal.Open(cas_input)
    out_cas_Open = gdal.Open(outputName)

    gdal_array.CopyDatasetInfo(temp_cas_Open, out_cas_Open)

    temp_cas_Open = out_cas_Open = None


if __name__ == '__main__':
    cas_input = ["D:/Pansharpening_Project/6/MY34_004209_158_0.cubeit.cub", "D:/Sincro/Pansharpening_Project/9/MY35_012250_189_0.cubeit.cub", "D:/Sincro/Pansharpening_Project/13/MY36_019052_194_3.cubeit.cub", "D:/Sincro/Pansharpening_Project/15/MY36_019318_220_0.cubeit.cub", "D:/Sincro/Pansharpening_Project/16/MY34_002260_337_0.cubeit.cub"]
    for n, i in enumerate(cas_input):
        CaSSIS_ColorSimulator_Main(i)
        print(f"{n+1}/{len(cas_input)} fatti!")