import numpy as np

from ..entities.initParamVariables import InitialCondition
from ..entities.modelConstants import ModelConstants
from typing import Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    # Important: classes are only imported when types are checked, not in production.
    from aquacrop.entities.inititalWaterContent import InitialWaterContent
    from aquacrop.entities.clockStruct import ClockStruct
    from aquacrop.entities.paramStruct import ParamStruct
    from aquacrop.entities.initParamVariables import InitialCondition


def read_model_initial_conditions(
    ParamStruct: "ParamStruct",
    ClockStruct: "ClockStruct",
    InitWC: "InitialWaterContent") -> Tuple["InitialCondition", "InitialCondition"]:
    """
    Function to set up initial model conditions

    Arguments:

        ParamStruct (ParamStruct):  Contains model paramaters

        ClockStruct (ClockStruct):  time paramaters

        InitWC (InitialWaterContent):  initial water content


    Returns:

        ParamStruct (ParamStruct):  updated ParamStruct object

        InitCond (InitialCondition):  containing initial model conditions/counters


    """

    ###################
    # creat initial condition class
    ###################

    InitCond = InitialCondition(len(ParamStruct.Soil.profile))

    # class_args = {key:value for key, value in InitCond_class.__dict__.items() if not key.startswith('__') and not callable(key)}
    # InitCond = InitCondStruct(**class_args)

    if ClockStruct.season_counter == -1:
        InitCond.z_root = 0.
        InitCond.cc0_adj = 0.

    elif ClockStruct.season_counter == 0:
        InitCond.z_root = ParamStruct.Seasonal_Crop_List[0].Zmin
        InitCond.cc0_adj = ParamStruct.Seasonal_Crop_List[0].CC0

    ##################
    # save field management
    ##################

    # Initial surface storage between any soil bunds
    if ClockStruct.season_counter == -1:
        # First day of simulation is in fallow period
        if (ParamStruct.FallowFieldMngt.bunds) and (
            float(ParamStruct.FallowFieldMngt.z_bund) > 0.001
        ):
            # Get initial storage between surface bunds
            InitCond.surface_storage = float(ParamStruct.FallowFieldMngt.bund_water)
            if InitCond.surface_storage > float(ParamStruct.FallowFieldMngt.z_bund):
                InitCond.surface_storage = float(ParamStruct.FallowFieldMngt.z_bund)
        else:
            # No surface bunds
            InitCond.surface_storage = 0

    elif ClockStruct.season_counter == 0:
        # First day of simulation is in first growing season
        # Get relevant field management structure parameters
        FieldMngtTmp = ParamStruct.FieldMngt
        if (FieldMngtTmp.bunds) and (float(FieldMngtTmp.z_bund) > 0.001):
            # Get initial storage between surface bunds
            InitCond.surface_storage = float(FieldMngtTmp.bund_water)
            if InitCond.surface_storage > float(FieldMngtTmp.z_bund):
                InitCond.surface_storage = float(FieldMngtTmp.z_bund)
        else:
            # No surface bunds
            InitCond.surface_storage = 0

    ############
    # watertable
    ############

    profile = ParamStruct.Soil.profile

    # Check for presence of groundwater table
    if ParamStruct.water_table == 0:  # No water table present
        # Set initial groundwater level to dummy value
        InitCond.z_gw = ModelConstants.NO_VALUE
        InitCond.wt_in_soil = False
        # Set adjusted field capacity to default field capacity
        InitCond.th_fc_Adj = profile.th_fc.values
    elif ParamStruct.water_table == 1:  # Water table is present
        # Set initial groundwater level
        InitCond.z_gw = float(ParamStruct.z_gw[ClockStruct.time_step_counter])
        # Find compartment mid-points
        zMid = profile.zMid
        # Check if water table is within modelled soil profile
        if InitCond.z_gw >= 0:
            idx = zMid[zMid >= InitCond.z_gw].index
            if idx.shape[0] == 0:
                InitCond.wt_in_soil = False
            else:
                InitCond.wt_in_soil = True
        else:
            InitCond.wt_in_soil = False

        # Adjust compartment field capacity
        compi = int(len(profile)) - 1
        thfcAdj = np.zeros(compi + 1)
        while compi >= 0:
            # get soil layer of compartment
            compdf = profile.loc[compi]
            if compdf.th_fc <= 0.1:
                Xmax = 1
            else:
                if compdf.th_fc >= 0.3:
                    Xmax = 2
                else:
                    pF = 2 + 0.3 * (compdf.th_fc - 0.1) / 0.2
                    Xmax = (np.exp(pF * np.log(10))) / 100

            if (InitCond.z_gw < 0) or ((InitCond.z_gw - zMid.iloc[compi]) >= Xmax):
                for ii in range(compi+1):
                    compdfii = profile.loc[ii]
                    thfcAdj[ii] = compdfii.th_fc

                compi = -1
            else:
                if compdf.th_fc >= compdf.th_s:
                    thfcAdj[compi] = compdf.th_fc
                else:
                    if zMid.iloc[compi] >= InitCond.z_gw:
                        thfcAdj[compi] = compdf.th_s
                    else:
                        dV = compdf.th_s - compdf.th_fc
                        dFC = (dV / (Xmax ** 2)) * ((zMid.iloc[compi] - (InitCond.z_gw - Xmax)) ** 2)
                        thfcAdj[compi] = compdf.th_fc + dFC

                compi = compi - 1

        # Store adjusted field capacity values
        InitCond.th_fc_Adj = np.round(thfcAdj, 3)

    profile["th_fc_Adj"] = np.round(InitCond.th_fc_Adj, 3)

    # create hydrology df to group by layer instead of compartment
    ParamStruct.Soil.Hydrology = profile.groupby("Layer").mean().drop(["dz", "dzsum"], axis=1)
    ParamStruct.Soil.Hydrology["dz"] = profile.groupby("Layer").sum().dz

    ###################
    # initial water contents
    ###################

    typestr = InitWC.wc_type
    methodstr = InitWC.method

    depth_layer = InitWC.depth_layer
    datapoints = InitWC.value

    values = np.zeros(len(datapoints))

    hydf = ParamStruct.Soil.Hydrology

    # Assign data
    if typestr == "Num":
        # Values are defined as numbers (m3/m3) so no calculation required
        depth_layer = np.array(depth_layer, dtype=float)
        values = np.array(datapoints, dtype=float)

    elif typestr == "Pct":
        # Values are defined as percentage of taw. Extract and assign value for
        # each soil layer based on calculated/input soil hydraulic properties
        depth_layer = np.array(depth_layer, dtype=float)
        datapoints = np.array(datapoints, dtype=float)

        for ii in range(len(values)):
            if methodstr == "Depth":
                depth = depth_layer[ii]
                value = datapoints[ii]

                # Find layer at specified depth
                if depth < profile.dzsum.iloc[-1]:
                    layer = profile.query(f"{depth}<dzsum").Layer.iloc[0]
                else:
                    layer = profile.Layer.iloc[-1]

                compdf = hydf.loc[layer]

                # Calculate moisture content at specified depth
                values[ii] = compdf.th_wp + ((value / 100) * (compdf.th_fc - compdf.th_wp))
            elif methodstr == "Layer":
                # Calculate moisture content at specified layer
                layer = depth_layer[ii]
                value = datapoints[ii]

                compdf = hydf.loc[layer]

                values[ii] = compdf.th_wp + ((value / 100) * (compdf.th_fc - compdf.th_wp))

    elif typestr == "Prop":
        # Values are specified as soil hydraulic properties (SAT, FC, or WP).
        # Extract and assign value for each soil layer

        for ii in range(len(values)):
            if methodstr == "Depth":
                # Find layer at specified depth
                depth = depth_layer[ii]
                value = datapoints[ii]

                # Find layer at specified depth
                if depth < profile.dzsum.iloc[-1]:
                    layer = profile.query(f"{depth}<dzsum").Layer.iloc[0]
                else:
                    layer = profile.Layer.iloc[-1]

                compdf = hydf.loc[layer]

                # Calculate moisture content at specified depth
                if value == "SAT":
                    values[ii] = compdf.th_s
                if value == "FC":
                    values[ii] = compdf.th_fc
                if value == "WP":
                    values[ii] = compdf.th_wp

            elif methodstr == "Layer":
                # Calculate moisture content at specified layer
                layer = depth_layer[ii]
                value = datapoints[ii]

                compdf = hydf.loc[layer]

                if value == "SAT":
                    values[ii] = compdf.th_s
                if value == "FC":
                    values[ii] = compdf.th_fc
                if value == "WP":
                    values[ii] = compdf.th_wp

    # Interpolate values to all soil compartments

    thini = np.zeros(int(profile.shape[0]))
    if methodstr == "Layer":
        for ii in range(len(values)):
            layer = depth_layer[ii]
            value = values[ii]

            idx = profile.query(f"Layer=={int(layer)}").index

            thini[idx] = value

        InitCond.th = thini

    elif methodstr == "Depth":
        depths = depth_layer

        # Add zero point
        if depths[0] > 0:
            depths = np.append([0], depths)
            values = np.append([values[0]], values)

        # Add end point (bottom of soil profile)
        if depths[-1] < ParamStruct.Soil.zSoil:
            depths = np.append(depths, [ParamStruct.Soil.zSoil])
            values = np.append(values, [values[-1]])

        # Find centroids of compartments
        SoilDepths = profile.dzsum.values
        comp_top = np.append([0], SoilDepths[:-1])
        comp_bot = SoilDepths
        comp_mid = (comp_top + comp_bot) / 2
        # Interpolate initial water contents to each compartment
        thini = np.interp(comp_mid, depths, values)
        InitCond.th = thini

    # If groundwater table is present and calculating water contents based on
    # field capacity, then reset value to account for possible changes in field
    # capacity caused by capillary rise effects
    if ParamStruct.water_table == 1:
        if (typestr == "Prop") and (datapoints[-1] == "FC"):
            InitCond.th = InitCond.th_fc_Adj

    # If groundwater table is present in soil profile then set all water
    # contents below the water table to saturation
    if InitCond.wt_in_soil is True:
        # Find compartment mid-points
        SoilDepths = profile.dzsum.values
        comp_top = np.append([0], SoilDepths[:-1])
        comp_bot = SoilDepths
        comp_mid = (comp_top + comp_bot) / 2
        idx = np.where(comp_mid >= InitCond.z_gw)[0][0]
        for ii in range(idx, len(profile)):
            layeri = profile.loc[ii].Layer
            InitCond.th[ii] = hydf.th_s.loc[layeri]

    InitCond.thini = InitCond.th

    ParamStruct.Soil.profile = profile
    ParamStruct.Soil.Hydrology = hydf

    return ParamStruct, InitCond

