from uuid import UUID
from pydantic import BaseModel, typing
from typing import List, Optional

from AvenirCommon.Logger import log
from AvenirCommon.Util import GBRange

from SpectrumCommon.Const.GB import *
from SpectrumCommon.Const.PJ import *


def getList_sex_singleAge_year(finalYrIdx, value = 0):
    return np.full((GB_Female + 1, GB_MaxSingleAges + 1, finalYrIdx + 1), value).tolist()

# def getList_5yrAge_sex_year(finalYrIdx, value = 0):
#     return np.full((GB_MAX_AGE + 1, GB_Female + 1, finalYrIdx + 1), value).tolist()

def getList_sex_5yrAge_year(finalYrIdx, value = 0):
    return np.full((GB_Female + 1, GB_MAX_AGE + 1, finalYrIdx + 1), value).tolist()

def getList_5yrAge_year(finalYrIdx, value = 0):
    return np.full((GB_MAX_AGE + 1, finalYrIdx + 1), value).tolist()

def getList_sex_year(finalYrIdx, value = 0):
    return np.full((GB_Female + 1, finalYrIdx + 1), value).tolist()

def getList_year(finalYrIdx, value = 0):
    return np.full(finalYrIdx + 1, value).tolist()

class calcParams(BaseModel):
    force: bool = False
    runImpact: bool = False
    runTBDynamicalModel: bool = False
    useLeapfrog: bool = False
    returnTags: list = []

class TaskMerge(BaseModel):
    id: UUID
    projID : str
    force : Optional[bool] = False

extra_dict = {
    "planType" : 0,
    GB_TB_COSTING_ON : False,
    GB_IHT_ON : False,
    GB_LIST_COSTING_ON : False,
    GB_GSKMode : False,
    "TBModel": ""
}

class createProjectionParams(BaseModel):
    fileName: str = "Test Projection"
    country: str = "BEN"
    firstYear: int = None
    finalYear: int = 2030
    modules: list = [GB_DP, GB_AM]
    userID: str = ""
    extra: dict = extra_dict

class NaomiSubnationalParams(BaseModel):
    country : str = "BEN"
    subnatRegion : str



class modifyProjectionParams(BaseModel):
    fileName: Optional[str] = None
    country: Optional[str] = None
    firstYear: Optional[int] = None
    finalYear: Optional[int] = None
    # modules: Optional[list] = None
    insertModules: Optional[list] = None
    removeModules: Optional[list] = None
    reloadModules: Optional[list] = None
    flatLinePastYears: Optional[bool] = False
    flatLineFutureYears: Optional[bool] = False
    TBSubnatRegion: Optional[str] = None
    TBImpactStatistical: Optional[bool] = None
    TBImpactDynamical: Optional[bool] = None
    appID: Optional[str] = None
    useLeapfrog: Optional[bool] = None


# class filterModules(BaseModel):
#    modvars: List[modvarParams] = [dict(modvarParams())]
class filterGetModule(BaseModel):
    module: str = "DP"
    modvars: list = ["<DP_FirstYear_V1>", "<DP_FinalYear_V1>"]


class filterGetParam(BaseModel):
    projID: str = "ceea11af94244344b4f4f9643e4c2b55"
    modules: List[filterGetModule] = [dict(filterGetModule())]


class filterSetModvar(BaseModel):
    tag: str = "<DP_FirstYear_V1>"
    value: typing.Any = 1970


class filterSetModule(BaseModel):
    module: str = "DP"
    modvars: List[filterSetModvar] = [dict(filterSetModvar())]


class filterSetParam(BaseModel):
    projID: str = "ceea11af94244344b4f4f9643e4c2b55"
    modules: List[filterSetModule] = [dict(filterSetModule())]


class LoginParams(BaseModel):
    username: str = "guest"
    password: str = "123456"


class datapackTreeParam(BaseModel):
    modSet: list = [1, 4]
    simple: bool = False
    configurations: bool = True
    configurationsKeyName: str = "config"


class datapackIndicatorsParam(BaseModel):
    modID: int = 1
    mstID: int = 1
    proj: str = ""

class citationParams(BaseModel):
    projID: str = ""
    clientVersion: str = ""

class datapackChartYears(BaseModel):
    firstYearIdx: int = 0
    finalYearIdx: int = 1
    min: int = 0
    max: int = 0


class datapackParam(BaseModel):
    chartYears: dict = dict(datapackChartYears())
    indicators: list = [dict(datapackIndicatorsParam())]
    language: int = 1

class getCountryParam(BaseModel):
    version: str = GBDatabaseVersion
    modID: int = GB_GB
    ignoreSampleCountry: bool = True

class getSubnationalParam(BaseModel):
    iso: int
    version: str = GBDatabaseVersion

class proximateDeterminantsParam(BaseModel):
    version: str = 'V2'

class missOpDataParam(BaseModel):
    countryID: str


#functions used for documenting information on modvars

def gb_shape(description : str, type, has_default=True, **kwargs):
    data_shape = {}
    data_shape['description'] = description
    data_shape['type'] = type
    data_shape['hasDefault'] = has_default
    data_shape.update(kwargs)
    return data_shape

gb_str_shape = lambda d, **kwargs : gb_shape(d, 'str', **kwargs)
gb_int_shape = lambda d, **kwargs : gb_shape(d, 'int', **kwargs)
gb_bool_shape = lambda d, **kwargs : gb_shape(d, 'bool', range='True, False', **kwargs)
gb_float_shape = lambda d, **kwargs : gb_shape(d, 'float', **kwargs)
gb_percent_shape  = lambda d, **kwargs : gb_shape(d, 'float', range='0-100', **kwargs)
gb_rate_shape  = lambda d, **kwargs : gb_shape(d, 'float', range='0.0-1.0', **kwargs)


gb_year_dim = 'year'
gb_year_indices = {'firstYear':0, 'finalYear':-1}

gb_sex_dim = 'sex'
gb_sex_indices = {'Both Sexes':GB_BothSexes, "Male":GB_Male, "Female":GB_Female}

gb_5y_age_dim = '5y_age'
gb_5y_age_indices = {
    'All Ages'      : GB_AllAges,
    'Age 0 to 4'    : GB_A0_4,
    'Age 5 to 9'    : GB_A5_9,
    'Age 10 to 14'  : GB_A10_14,
    'Age 15 to 19'  : GB_A15_19,
    'Age 20 to 24'  : GB_A20_24,
    'Age 25 to 29'  : GB_A25_29,
    'Age 30 to 34'  : GB_A30_34,
    'Age 35 to 39'  : GB_A35_39,
    'Age 40 to 44'  : GB_A40_44,
    'Age 45 to 49'  : GB_A45_49,
    'Age 50 to 54'  : GB_A50_54,
    'Age 55 to 59'  : GB_A55_59,
    'Age 60 to 64'  : GB_A60_64,
    'Age 65 to 69'  : GB_A65_69,
    'Age 70 to 74'  : GB_A70_74,
    'Age 75 to 79'  : GB_A75_79,
    'Age 80+'       : GB_A80_Up,
}

gb_5y_age_adult_dim = '5y_age_adult'
gb_5y_age_adult_indices = {
    'Age 15 to 19'  : GB_A15_19,
    'Age 20 to 24'  : GB_A20_24,
    'Age 25 to 29'  : GB_A25_29,
    'Age 30 to 34'  : GB_A30_34,
    'Age 35 to 39'  : GB_A35_39,
    'Age 40 to 44'  : GB_A40_44,
    'Age 45 to 49'  : GB_A45_49,
}

gb_single_age_dim = 'single_age'
gb_single_age_indices = {
    ('Age ' + str(a)) : a for a in GBRange(0, 80)
}

gb_single_age_child_dim = 'single_age_child'
gb_single_age_child_indices = {
    ('Age ' + str(a)) : a for a in GBRange(0, 14)
}

gb_single_age_child_17_dim = 'single_age_child_17'
gb_single_age_child_17_indices = {
    ('Age ' + str(a)) : a for a in GBRange(0, 17)
}

gb_single_age_child_4_dim = 'single_age_child_4'
gb_single_age_child_4_indices = {
    ('Age ' + str(a)) : a for a in GBRange(0, 4)
}

gb_year_shape  =  lambda d, **kwargs  : gb_shape(d,'1D float', dim=[gb_year_dim], indices=[gb_year_indices], **kwargs)
gb_sex_year_shape  =  lambda d, **kwargs  : gb_shape(d, '2D float',
                                                     dim=[gb_sex_dim, gb_year_dim],
                                                     indices=[gb_sex_indices, gb_year_indices], **kwargs)
gb_sex_single_age_year_shape  =  lambda d, **kwargs  : gb_shape(d, '3D float',
                                                     dim=[gb_sex_dim, gb_single_age_dim, gb_year_dim],
                                                     indices=[gb_sex_indices, gb_single_age_indices, gb_year_indices], **kwargs)
gb_sex_5y_age_year_shape  =  lambda d, **kwargs  : gb_shape(d, '3D float',
                                                     dim=[gb_sex_dim, gb_5y_age_dim, gb_year_dim],
                                                     indices=[gb_sex_indices, gb_5y_age_indices, gb_year_indices], **kwargs)
gb_5y_age_year_shape  =  lambda d, **kwargs  : gb_shape(d, '2D float',
                                                     dim=[gb_5y_age_dim, gb_year_dim],
                                                     indices=[gb_5y_age_indices, gb_year_indices], **kwargs)

def gb_reshape_years(mv_piece, meta_mv, original_first_year, original_final_year, new_first_year, new_final_year):

    # Convert lists to numpy arrays if they aren't already numpy arrays.
    is_list = type(mv_piece)==list
    if is_list:
        mv_piece = np.array(mv_piece)

    new_num_years = new_final_year - new_first_year + 1

    # Determine the range of years to copy from the original array to the new array.
    copy_start_year = max(original_first_year, new_first_year)
    copy_end_year = min(original_final_year, new_final_year)

    original_start_index = copy_start_year - original_first_year # index of original year dimension to start copying data
    original_end_index = copy_end_year - original_first_year # index of original year dimension to stop copying data.
    new_start_index = copy_start_year - new_first_year # index of new year dimension to start pasting data
    new_end_index = copy_end_year - new_first_year # index of new year dimension to stop pasting data

    # Identify the year dimension.
    year_dim = meta_mv['dim'].index('year')

    # Create a new array with the new number of years.
    new_shape = list(mv_piece.shape) # dimensions of array (shape) as list
    new_shape[year_dim] = new_num_years # update year dimension to new number of years
    new_mv_piece = np.zeros(new_shape, dtype=mv_piece.dtype) # create new mv_piece based on shape with updated number of years and same data type as original

    # Copy the relevant data from the original array to the new array. ... = include all dimensions up to the point of the year dimension.
    new_mv_piece[..., new_start_index:new_end_index+1] = mv_piece[..., original_start_index:original_end_index+1]

    # Fill in new years beyond the original final year with the value of the original final year
    if new_final_year > original_final_year:
        final_year_value = mv_piece.take(indices=original_end_index, axis=year_dim) # Grab elements at original_end_index along the year dimension.

        for t in range(new_end_index+1, new_num_years):
            new_mv_piece[..., t] = final_year_value

    # Turn numpy type back into the list or set the new mv_piece to the original otherwise.
    mv_piece = new_mv_piece.tolist() if is_list else new_mv_piece

    return mv_piece

# original
# def gb_reshape_years(mv_piece, meta_mv, original_first_year, original_final_year, new_first_year, new_final_year):

#     first_year_offset = new_first_year - original_first_year
#     final_year_offset = new_final_year - original_final_year
#     final_year_index =  original_final_year - original_first_year

#     #temp convert list to numpy type
#     is_list = type(mv_piece)==list
#     if is_list:
#         mv_piece = np.array(mv_piece)

#     axis = 0
#     pad_by_dim = []
#     for i, e in enumerate(meta_mv['dim']):
#         if e == gb_year_dim:
#             axis = i
#             pad_by_dim.append((0, final_year_offset))
#         else:
#             pad_by_dim.append((0, 0))


#     #extending final year
#     if final_year_offset>0:

#         #Reshapes an N dimensional array for years
#         pad_start = [0]*len(mv_piece.shape)
#         pad_end   = [0]*len(mv_piece.shape)
#         pad_end[axis] = final_year_offset
#         mv_piece = np.pad(mv_piece, pad_by_dim, mode='edge')


#     #extending first year or decreasing final year
#     if (first_year_offset>0) or  (new_final_year<=original_final_year):
#         mv_piece = mv_piece.take(indices=range(first_year_offset, final_year_index+1), axis=axis)

#     #turn numpy type back into the list
#     if is_list:
#         mv_piece = mv_piece.tolist()

#     return mv_piece

def gb_update_years_inner(mv_piece, meta_mv, original_first_year, original_final_year, new_first_year, new_final_year):
    # If a year dimension was specified for the shape, pass the piece of the MV off to get reshaped based on the new years.
    if gb_year_dim in meta_mv.get('dim', []):
        if len(mv_piece) > 0:
            mv_piece = gb_reshape_years(mv_piece,
                                        meta_mv,
                                        original_first_year,
                                        original_final_year,
                                        new_first_year,
                                        new_final_year)
            return mv_piece

        # If contingent=True, we want to ignore processing the piece for whatever reason, so just return it as-is.
        elif ('contingent' in meta_mv) and (meta_mv['contingent']):
            return mv_piece

        # Otherwise, the MV should be fixed in the shapes file.
        else:
            raise Exception('Year dimension specified in MV shapes file, but MV is of length 0. Either correct the problem, or add contingent=True to the MV shape.')

    elif meta_mv['type'] == 'dict':
        if 'innerShape' in meta_mv:
            for key in meta_mv['innerShape']:
                inner_meta = meta_mv['innerShape'][key]
                # Some MVs' parts may or may not be present. If they are missing and should always be there, then raise an exception.
                # If they are missing but it's ok, just skip over processing the key.
                if key in mv_piece :
                    # print(key)
                    try:
                        inner_mv = mv_piece[key]
                        mv_piece[key] = gb_update_years_inner(inner_mv, inner_meta, original_first_year, original_final_year, new_first_year, new_final_year)
                    except Exception as e:
                        raise Exception(f"Error processing key '{key}' in MV piece: {e}")


                elif ('contingent' in meta_mv) and (meta_mv['contingent']):
                    return mv_piece

                else:
                    continue # skip processing this key if it is not present in the MV piece.
                    # raise Exception('Key \'' + key + '\' specified in MV shapes file was not found in the MV. Either correct the problem, or add contingent=True to the MV shape.')
            return mv_piece
        else:
            if ('contingent' in meta_mv) and (meta_mv['contingent']):
                return mv_piece

            else:
                pass #no innershape specified so assumption is years do not exist in this MV piece.
                return mv_piece
            #     raise Exception('innerShape was not specified for dict type. Either correct the problem, or add contingent=True to the MV shape.')

    #???
    elif((meta_mv['type'] == 'dictlist') or (meta_mv['type'] == '1D dict')):
        for i, piece in enumerate(mv_piece):

            if 'type' in meta_mv['innerShape']:
                inner_meta = meta_mv['innerShape']
            else:
                inner_meta = {
                    'innerShape': meta_mv['innerShape'],
                    'type': 'dict',  # Default to dict if no type is specified.
                }
                if 'contingent' in meta_mv:
                    inner_meta['contingent'] = meta_mv['contingent']
            mv_piece[i] = gb_update_years_inner(piece, inner_meta, original_first_year, original_final_year, new_first_year, new_final_year)
        return mv_piece

    # Otherwise, no years to update; return same MV that was passed in.
    else:
        return mv_piece

def gb_update_years(projection, params, shapes):
    first_year_offset = params.firstYear - projection[PJN_FirstYearTag]
    final_year_offset = params.finalYear - projection[PJN_FinalYearTag]
    final_year_index =  projection[PJN_FinalYearTag] - projection[PJN_FirstYearTag]


    for tag in shapes:
        print(tag)
        if tag == '<IS_FacilityRecords_V2>':
            pass
        # if tag == '<CS_SBPrevByInterResults_V1>':
        #     pass
        if (tag not in projection):
            pass
            continue  # Skip if the tag is not present in the projection.

        try:
            meta_mv = shapes[tag]
            projection[tag] = gb_update_years_inner(projection[tag],
                                                    meta_mv,
                                                    projection[PJN_FirstYearTag],
                                                    projection[PJN_FinalYearTag],
                                                    params.firstYear,
                                                    params.finalYear)
        except Exception as e:
            log(f"Error updating years for tag '{tag}': {e}")
            raise Exception(f"Error updating years for tag '{tag}': {e}")
    return None

# By no means finished, but goal is to eventually spit out a shapes dict for a module given it's MVs.  Currently just prints out the MVs with their
def gb_generate_shapes_skeleton(MVs):
    for key, value in MVs.items():
        var_type = str(type(value).__name__)

        dimensions = 0
        if var_type == 'list':
            def get_list_dimensions(lst):
                if not isinstance(lst, list):
                    return 0
                if not lst: # An empty list is still 1-dimensional
                    return 1
                return 1 + get_list_dimensions(lst[0])

            def get_list_data_type(lst):
                # Base case: if lst is not a list, return its type
                if not isinstance(lst, list):
                    return type(lst).__name__
                # Recursive case: drill down into the list
                if lst: # check if the list is not empty
                    return get_list_data_type(lst[0])
                else:
                    return 'Empty list' # You can decide how to handle empty lists

            dimensions = get_list_dimensions(value)
            data_type = str(get_list_data_type(value))
            var_type += ", " + str(dimensions) + " dimensions, data type: " + data_type

        log(key + ': ' + var_type)
