'''
/***************************************************************************
Name		: under_catch_worker.py
Description : BASIICS Under Catch Correct Stations for Fews_Tools Plugin
copyright   : (C) 2020-2026 by FEWS
email       : dhackman@contractor.usgs.gov
Created     : 02/03/2026 dhackman
Modified    :
 ***************************************************************************/

/***************************************************************************
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 ***************************************************************************/
'''
import os
import numpy as np

from PyQt5 import QtCore
from qgis.core import QgsMessageLog, Qgis

from fews_tools import fews_tools_config as config
from fews_tools.utilities import basiics_batch_utilities as b_util
from fews_tools.utilities import geoclim_gdal_utilities as g_util
from fews_tools.utilities import geoclim_utilities as util
from fews_tools.utilities import geoclim_qgs_utilities as qgs_util

class UnderCatchWorker(QtCore.QObject):
    '''
    Class for Under Catch Correcting Station files
    '''
    def __init__(self, batch_dic, wrksp_setup):
        '''
        init
        params(dic) - batch_dic
        params(dic) - wrksp_setup
        '''
        QtCore.QObject.__init__(self)
        self.bat_dic = batch_dic
        self.wrksp_setup = wrksp_setup
        self.stn_dic = None
        self.killed = False
        self.step = 10.0
        self.curr_progress = 0.0
        self.out_csv_settings = {
            'delimiter': ',',
            'stn_missing_val': -9999,
            'csv_long_col': 0,
            'csv_lat_col': 0,
            'csv_year_col': 0,
            'csv_hdr_row': 1,
            'csv_stn_id_col': 0,
            'csv_beg_period_col': 0,
            'csv_end_period_col': 0,
        }
        if self.bat_dic['update_existing_corrected_file']:
            # Set values for out_csv to ones set for corrected file
            self.out_csv_settings['delimiter'] =\
                self.bat_dic['delimiter2']
            self.out_csv_settings['stn_missing_val'] =\
                self.bat_dic['stn_missing_val2']
            self.out_csv_settings['csv_long_col'] =\
                self.bat_dic['csv_long_col2']
            self.out_csv_settings['csv_lat_col'] =\
                self.bat_dic['csv_lat_col2']
            self.out_csv_settings['csv_year_col'] =\
                self.bat_dic['csv_year_col2']
            self.out_csv_settings['csv_hdr_row'] =\
                self.bat_dic['csv_hdr_row2']
            self.out_csv_settings['csv_stn_id_col'] =\
                self.bat_dic['csv_stn_id_col2']
            self.out_csv_settings['csv_beg_period_col'] =\
                self.bat_dic['csv_beg_period_col2']
            self.out_csv_settings['csv_end_period_col'] =\
                self.bat_dic['csv_end_period_col2']
            self.out_csv_settings['stn_file_col_list'] =\
                self.bat_dic['stn_file_col_list2']
            self.input_corrected_csv_data_list = b_util.read_csv_file(
                self.bat_dic['corrected_station_filename'],
                self.bat_dic['delimiter2'], self.bat_dic['csv_hdr_row2'])
        else:
            # Set values to same as the input_csv since no existing file,
            # except still use delimeter/header row/missing value from
            # output section
            self.out_csv_settings['delimiter'] =\
                self.bat_dic['delimiter2']
            self.out_csv_settings['stn_missing_val'] =\
                self.bat_dic['stn_missing_val2']
            self.out_csv_settings['csv_long_col'] =\
                self.bat_dic['csv_long_col']
            self.out_csv_settings['csv_lat_col'] =\
                self.bat_dic['csv_lat_col']
            self.out_csv_settings['csv_year_col'] =\
                self.bat_dic['csv_year_col']
            self.out_csv_settings['csv_hdr_row'] =\
                self.bat_dic['csv_hdr_row2']
            self.out_csv_settings['csv_stn_id_col'] =\
                self.bat_dic['csv_stn_id_col']
            self.out_csv_settings['csv_beg_period_col'] =\
                self.bat_dic['csv_beg_period_col']
            self.out_csv_settings['csv_end_period_col'] =\
                self.bat_dic['csv_end_period_col']
            self.out_csv_settings['stn_file_col_list'] =\
                self.bat_dic['stn_file_col_list']
            self.input_corrected_csv_data_list = None

    def run(self):
        '''
        run method
        '''
        ret_tuple = None
        try:
            QgsMessageLog.logMessage(
                'Beginning under catch correction of stations',
                level=Qgis.Info)

            # Always read in the input csv
            input_csv_data_list = b_util.read_csv_file(
                    self.bat_dic['station_filename'],
                    self.bat_dic['delimiter'], self.bat_dic['csv_hdr_row'])
            self.update_progress()
            if self.killed is True:
                raise KeyboardInterrupt

            output_corrected_csv_data_list = []
            # Add Header Row, either the input file headers or file
            # we are updating's headers
            self.headers_line = self.out_csv_settings['stn_file_col_list']
            for idx in range(self.bat_dic['csv_beg_period_col'] - 1,
                             self.bat_dic['csv_end_period_col']):
                uf_col = 'uf_' +\
                    self.out_csv_settings['stn_file_col_list'][idx]
                if uf_col not in self.headers_line:
                    self.headers_line.append(uf_col)
            output_corrected_csv_data_list.append(self.headers_line)

            # Build col/month dict based off each valid year for lookup later
            self.col_month_year_dict = self.build_col_month_year_dict()
            self.update_progress()
            if self.killed is True:
                raise KeyboardInterrupt

            # Build corrfactor data dict for lookup later
            self.corrfactor_data_dict = \
                self.build_corrfactor_lookup_dict()
            self.update_progress()
            if self.killed is True:
                raise KeyboardInterrupt

            self.step = 40.0/len(input_csv_data_list)
            for input_line in input_csv_data_list:
                self.update_progress()
                if self.killed is True:
                    raise KeyboardInterrupt
                year = int(input_line[self.bat_dic['csv_year_col'] - 1])
                col_month_list = self.col_month_year_dict.get(year, None)
                if not col_month_list:
                    # Not in range that we are processing, so skip this line
                    continue

                if self.input_corrected_csv_data_list:
                    # Get same row from existing corrected file if it exists
                    # POP row off the input_corrected_csv_data_list if it
                    # matches so we know which rows are processed or not
                    input_corrected_csv_row = \
                        self.get_existing_corrected_row(input_line)
                else:
                    input_corrected_csv_row = None

                # for each tuple in list loop through and calculate the
                # corrected value based on lat/long and output a new
                # output tuple with final output
                output_line = self.calc_corrected_row(
                    input_line, col_month_list, input_corrected_csv_row)

                # append new line to output_corrected_csv_data_list
                output_corrected_csv_data_list.append(output_line)

            self.step = 10.0
            self.update_progress()
            if self.killed is True:
                raise KeyboardInterrupt
            # add any leftover input_corrected_csv_data_list rows
            if self.input_corrected_csv_data_list:
                output_corrected_csv_data_list.extend(self.input_corrected_csv_data_list)

            # organize the list based on year and station ID
            output_corrected_csv_data_list =\
                [output_corrected_csv_data_list[0]] + sorted(
                    output_corrected_csv_data_list[1:], key=lambda row: (
                row[self.out_csv_settings['csv_year_col'] - 1],
                row[self.out_csv_settings['csv_stn_id_col'] - 1]))

            self.update_progress()
            if self.killed is True:
                raise KeyboardInterrupt
            # Write out list to the csv
            with open(self.bat_dic['corrected_station_filename'],
                      'w') as corr_file:
                for line in output_corrected_csv_data_list:
                    output = ''
                    for col in line:
                        output += str(col) + self.out_csv_settings['delimiter']
                    output += '\n'
                    corr_file.write(output)

            if self.killed is False:
                self.progress.emit(100)
                ret_tuple = (0, "Under Catch Correct Stations Complete")
        # exit with appropriate message on killed (KeyboardInterrupt)
        except KeyboardInterrupt:
            self.progress.emit(0)
            ret_tuple = (0, u"Under Catch Correct Stations aborted by user")
        # forward any execeptions upstream
        except BaseException as exc:
            print(exc)
            self.error.emit(
                exc, u"Unspecified error in Under Catch Correct Stations")
        self.finished.emit(ret_tuple)

    def calc_corrected_row(self, input_line, input_col_month_list, input_corr_row):
        '''
        Will run the calculations for each column in date range and add it to
        output file 

        args:
            input_line: list(strings) - input line from the csv
            input_col_month_list: list(tuples) - col index and corrfactor file
                for index
            input_corr_row: list(strings) - matching input line from the 
                corrected csv if it exists
        returns:
            output_line: list(strings) - list of all columns for output file
                with the corrected columns generated
        '''
        output_line = []
        corr_factors_list = []
        col_idx_dif = self.bat_dic['csv_beg_period_col'] -\
            self.out_csv_settings['csv_beg_period_col']
        if input_corr_row:
            # Get lat and long, default to using the input_corr_row since that
            # is what is in the existing output
            lat = float(
                input_corr_row[self.out_csv_settings['csv_lat_col'] - 1])
            long = float(
                input_corr_row[self.out_csv_settings['csv_long_col'] - 1])
            # Set all non period column values, since there is an input
            # corr row it is easy and we can just copy them over in order,
            # since already set
            for idx in range(0, self.out_csv_settings['csv_beg_period_col'] - 1):
                output_line.append(input_corr_row[idx])
        else:
            lat = float(input_line[self.bat_dic['csv_lat_col'] - 1])
            long = float(input_line[self.bat_dic['csv_long_col'] - 1])
            # Set all non period column values, by matching input indexes to
            # output file indexes, since there is not already an input corr
            # row for this entry
            index_val_list = [None] * \
                (self.out_csv_settings['csv_beg_period_col'] - 1)

            # Set indexes we know to correct location in list
            index_val_list[self.out_csv_settings['csv_stn_id_col'] - 1] =\
                input_line[self.bat_dic['csv_stn_id_col'] - 1]
            index_val_list[self.out_csv_settings['csv_year_col'] - 1] =\
                input_line[self.bat_dic['csv_year_col'] - 1]
            index_val_list[self.out_csv_settings['csv_long_col'] - 1] =\
                input_line[self.bat_dic['csv_long_col'] - 1]
            index_val_list[self.out_csv_settings['csv_lat_col'] - 1] =\
                input_line[self.bat_dic['csv_lat_col'] - 1]

            # Loop through for unknown columns that might still exist
            for i in index_val_list:
                if i == None:
                    idx = index_val_list.index(None) # Return first missing index
                    # Loop indexes to get first value not in output list
                    # assume that is correct value for that column
                    for idx2 in range(0, self.bat_dic['csv_beg_period_col'] - 1):
                        if input_line[idx2] not in index_val_list:
                            index_val_list[idx] = input_line[idx2]
                            break

            # If any nones left in list replace None with ''
            output_line = ['' if x == None else x for x in index_val_list]

        # calculate x and y value based on the lat and longitude
        for in_tuple in input_col_month_list:
            # get value at index (-1) and get dict key's value tuple
            # array and metadata for .tif file in tuple
            in_idx = in_tuple[0] - 1
            out_idx = in_idx - col_idx_dif
            corrfactor_file = in_tuple[1]
            initial_val = float(input_line[in_idx])
            if corrfactor_file == 'Outside Range':
                corr_factor = 'N/A'
                if input_corr_row:
                    # leave it the same as the input line
                    output_value = float(input_corr_row[out_idx])
                    # Get the input from input line for index of uf period
                    if ('uf_' + self.headers_line[out_idx]) in self.headers_line:
                        out_uf_idx = self.headers_line.index('uf_' + self.headers_line[out_idx])
                        if out_uf_idx < len(input_corr_row):
                            corr_factor = input_corr_row[out_uf_idx]
                else:
                    if initial_val == self.bat_dic['stn_missing_val']:
                        # Or set as stn missing val if input was
                        output_value = self.out_csv_settings['stn_missing_val']
                    else:
                        # Or set as input if not a missing val
                        output_value = initial_val
            else:
                # Calc value based on long and lat using corrfactor file
                output_value, corr_factor = self.calc_corr_val_from_stn_loc(
                    initial_val, corrfactor_file, lat, long)

            corr_factors_list.append(str(corr_factor))
            output_line.append(str(output_value))
        
        output_line.extend(corr_factors_list)

        return output_line

    def calc_corr_val_from_stn_loc(self, initial_val, corrfactor_file, lat, long):
        '''
        This will take in the initial value, corrfactor file name, and station
        info and calulate the final corrected output value for that station,
        year, and period index

        args:
            initial_val: int/float - value for station for year and period
            corrfactor_file: string - basename of corrfactor file for that
                periods month, can be used to lookup in the 
                self.col_month_year_dict
            lat: float - station's lattidude
            long: float - station's longitude
        returns:
            output_line: float - final value after applying undercatch
            corr_factor: float - the corr factor being appied
        '''
        output_val = self.out_csv_settings['stn_missing_val']
        corrfactor_entry = self.corrfactor_data_dict.get(
            corrfactor_file, None)
        if corrfactor_entry:
            grid_arr = corrfactor_entry[0]
            params = corrfactor_entry[1]
            pt_col = int((long - params[0].xMinimum()) / params[3])
            pt_row = int((params[0].yMaximum() - lat) / params[3])
            corr_factor = grid_arr[pt_row, pt_col]
            if initial_val != self.bat_dic['stn_missing_val']:
                output_val = float(initial_val * corr_factor)

        return output_val, corr_factor

    def get_existing_corrected_row(self, input_line):
        '''
        This will search the self.input_corrected_csv_data_list if it is not
        None and will pop a matching row off of that list and return it for
        future use in the process

        args:
            input_line: list(strings) - input line from the csv

        return:
            matching_input_corr_csv_row: list(strings) - if it exists
        '''
        matching_input_corr_csv_row = None

        idx = -1
        for index, sublist in enumerate(self.input_corrected_csv_data_list):
            if (input_line[self.bat_dic['csv_stn_id_col'] - 1] ==\
                    sublist[self.out_csv_settings['csv_stn_id_col'] - 1] and 
                    input_line[self.bat_dic['csv_year_col'] - 1] ==\
                    sublist[self.out_csv_settings['csv_year_col'] - 1]):
                # Both station id and year match then that is line we need
                idx = index  # Set the index of the list

        if idx != -1:
            matching_input_corr_csv_row = \
                self.input_corrected_csv_data_list.pop(idx)

        return matching_input_corr_csv_row

    def build_corrfactor_lookup_dict(self):
        '''
        Will build a dictionary for each months corrfactor file by extcrating
        an array of each 
        
        returns:
            corrfactor_data_dict: dict(numpy arr, file params) - numpy arrays for each
                monthly corrfactor file
        '''
        corrfactor_data_dict = {}
        for file in self.bat_dic['undercatch_fields_file_list']:
            _, _, _, _, data_type = g_util.get_geotiff_info(file)
            np_data_type = g_util.TYPE_DIC[data_type]["NP"]
            data_arr = g_util.extract_raster_array(file, np_data_type)
            params = qgs_util.extract_raster_file_params(file)
            corrfactor_data_dict[os.path.basename(file)] = (data_arr, params)

        return corrfactor_data_dict

    def build_col_month_year_dict(self):
        '''
        returns:
            col_month_list: dict(tuples) - columns to process paired with
                month string they fall under
        '''
        col_month_list = {}
        interval_dic = util.get_interval_dic(self.bat_dic['periodicity'])

        for year in range(int(self.bat_dic['start_year']),
                         int(self.bat_dic['end_year']) + 1):
            start_year = False
            end_year = False
            # These are not if else's in case they have same start and end year
            if year == int(self.bat_dic['start_year']):
                start_year = True
            if year == int(self.bat_dic['end_year']):
                end_year = True

            col_month_list[year] = self.get_month_for_valid_periods(
                interval_dic, start_year, end_year)

        return col_month_list

    def get_month_for_valid_periods(self, period_dict, start_year,
                                    end_year):
        '''
        This will use the periodicity dict to generate a list of tuples
        
        returns:
            col_month_list: list(tuples) - tuple is the column and month for
                that column
        '''
        col_month_list = []
        if start_year:
            start_column = \
                list(period_dict.keys()).index(self.bat_dic['start_period']) +\
                self.bat_dic['csv_beg_period_col']
        else:
            start_column = self.bat_dic['csv_beg_period_col']

        if end_year:
            end_column = \
                list(period_dict.keys()).index(self.bat_dic['end_period']) +\
                self.bat_dic['csv_beg_period_col']
        else:
            end_column = self.bat_dic['csv_end_period_col']

        periods = len(list(period_dict.keys()))
        periods_per_month = int(periods/12)
        month_dic_list = list(config.MONTH_2_NUM_DIC.items())
        for period_idx in range(self.bat_dic['csv_beg_period_col'],
                                self.bat_dic['csv_end_period_col'] + 1):
            if period_idx in range(start_column, end_column + 1):
                period = period_idx - self.bat_dic['csv_beg_period_col']
                month = int(period/periods_per_month)
                month_str = self.bat_dic['input_prefix'] +\
                      month_dic_list[month][1]['PER_YEAR'] +\
                        config.TIF_SUFFIX
            else:
                # Invalid period, outside of run range
                month_str = 'Outside Range'
            col_month_list.append((period_idx, month_str))
        return col_month_list

    def kill(self):
        '''
        Kill method
        '''
        self.killed = True

    def update_progress(self):
        '''
        Helper for progress bar updates
        '''
        self.curr_progress += self.step
        if self.curr_progress > 90:
            self.curr_progress = 90
        self.progress.emit(int(self.curr_progress))

    finished = QtCore.pyqtSignal(object)

    error = QtCore.pyqtSignal(Exception, str)

    progress = QtCore.pyqtSignal(int)
