import matplotlib.pyplot as plt
import datetime
import numpy as np
from typing import Optional, List, Dict

from snowpat.icsv import iCSV2DTimeseries
from snowpat.snowpackreader import SnowpackReader, Snowpack

from .Utils import show_figure


DEFAULT_FIELD_CODE_MAP = {
    "layer_index": "0501",
    "z": "0501",
    "thickness": "0501",
    "density": "0502",
    "temperature": "0503",
    "volumetric_water_content": "0506",
    "liquid_water_content": "0506",
    "dendricity": "0508",
    "sphericity": "0509",
    "sph": "0509",
    "coordination_number": "0510",
    "bond_size": "0511",
    "grain_size": "0512",
    "grain_type": "0513",
    "volumetric_ice_content": "0515",
    "ice_volume_fraction": "0515",
    "air_volume_fraction": "0516",
    "stress": "0517",
    "viscosity": "0518",
    "soil_volume_fraction": "0519",
    "temperature_gradient": "0520",
    "thermal_conductivity": "0521",
    "absorbed_shortwave_radiation": "0522",
    "viscous_deformation_rate": "0523",
    "stability_index_sdef": "0531",
    "natural_stability_index": "0532",
    "stability_index_sk38": "0533",
    "hand_hardness": "0534",
    "optical_equivalent_grain_size": "0535",
    "snow_shear_strength": "0601",
    "grain_size_difference": "0602",
    "hardness_difference": "0603",
    "ssi": "0604",
    "inverse_texture_index": "0605",
    "critical_cut_length": "0606",
}


def build_field_code_map(icsv_file: iCSV2DTimeseries, field_code_map: Optional[Dict[str, str]] = None) -> Dict[str, str]:
    if not isinstance(icsv_file, iCSV2DTimeseries):
        raise ValueError(f"An iCSV2DTimeseries object has to be provided, got {type(icsv_file).__name__}")

    merged_map = dict(DEFAULT_FIELD_CODE_MAP)
    if field_code_map:
        merged_map.update({key.lower(): value for key, value in field_code_map.items()})

    fallback_code_counter = 7000
    resolved_map = {}
    used_codes = set()
    for field_name in icsv_file.fields.fields:
        normalized = field_name.lower()
        code = merged_map.get(normalized)
        if code is None or code in used_codes:
            while True:
                code = f"{fallback_code_counter:04d}"
                fallback_code_counter += 1
                if code not in used_codes:
                    break
        resolved_map[field_name] = code
        used_codes.add(code)
    return resolved_map


def to_snowpack(icsv_file: iCSV2DTimeseries, field_code_map: Optional[Dict[str, str]] = None) -> SnowpackReader:
    if not isinstance(icsv_file, iCSV2DTimeseries):
        raise ValueError(f"An iCSV2DTimeseries object has to be provided, got {type(icsv_file).__name__}")

    snowpack = SnowpackReader.__new__(SnowpackReader)
    snowpack.filename = icsv_file.filename
    snowpack.metadata = getattr(icsv_file, "metadata", {})
    snowpack._old_hardness = False
    snowpack.DataCodes = {}
    snowpack.units = {}
    snowpack.data = {}
    snowpack._field_to_code = build_field_code_map(icsv_file, field_code_map)

    _set_metadata_from_fields(snowpack, icsv_file)
    _build_profiles(snowpack, icsv_file)
    snowpack.NamesToCodes = {v: k for k, v in snowpack.DataCodes.items()}
    return snowpack


def map_field_names_to_var_codes(field_names: Optional[List[str]], icsv_file: iCSV2DTimeseries, field_code_map: Dict[str, str]) -> List[str]:
    if field_names is None:
        field_names = [field_name for field_name in icsv_file.fields.fields if field_name != "layer_index"]

    available_fields = icsv_file.fields.fields
    for field_name in field_names:
        if field_name not in available_fields:
            raise ValueError(f"Invalid field name: {field_name}")

    return [field_code_map[field_name] for field_name in field_names]


def map_field_kwargs_to_var_codes(kwargs: Dict, field_code_map: Dict[str, str]) -> Dict:
    mapped_kwargs = dict(kwargs)
    code_field_map = {code: field_name for field_name, code in field_code_map.items()}

    for key in ["cmap", "norm", "cbar_label", "vmin", "vmax", "subtitle"]:
        value = mapped_kwargs.get(key)
        if isinstance(value, dict):
            mapped_kwargs[key] = {
                field_code_map.get(field_name, field_name): field_value
                for field_name, field_value in value.items()
            }

    adjust_data = mapped_kwargs.get("adjust_data")
    if adjust_data is not None:
        def wrapped(var_code, values):
            field_name = code_field_map.get(var_code, var_code)
            return adjust_data(field_name, values)

        mapped_kwargs["adjust_data"] = wrapped

    return mapped_kwargs


def _set_metadata_from_fields(snowpack: SnowpackReader, icsv_file: iCSV2DTimeseries):
    units = []
    long_names = []
    if getattr(icsv_file.fields, "units", None):
        units = icsv_file.fields.units.split("|")
    if getattr(icsv_file.fields, "long_name", None):
        long_names = icsv_file.fields.long_name.split("|")

    for idx, field_name in enumerate(icsv_file.fields.fields):
        code = snowpack._field_to_code[field_name]
        unit = units[idx].strip() if idx < len(units) and units[idx].strip() != "-" else "b.E."
        long_name = long_names[idx].strip() if idx < len(long_names) and long_names[idx].strip() != "-" else field_name
        snowpack.units[code] = unit
        snowpack.DataCodes[code] = long_name


def _build_profiles(snowpack: SnowpackReader, icsv_file: iCSV2DTimeseries):
    for date in icsv_file.dates:
        df = icsv_file.data[date]
        snowpack.data[date] = _profile_from_dataframe(df, snowpack._field_to_code)


def _profile_from_dataframe(df, field_code_map: Dict[str, str]) -> Snowpack:
    profile = Snowpack()
    boundaries = _extract_layer_boundaries(df)
    profile.set_param("0501", boundaries, len(boundaries))

    n_layers = len(boundaries) - 1
    for field_name in df.columns:
        if field_name == "layer_index":
            continue
        code = field_code_map[field_name]
        values = _prepare_layer_values(df[field_name].to_numpy(), n_layers)
        profile.set_param(code, values, len(boundaries))

    profile._parse_data(False)
    profile._parsed = True
    return profile


def _extract_layer_boundaries(df) -> np.ndarray:
    if "Z" in df.columns:
        z_values = df["Z"].to_numpy(dtype=float)
        z_values = np.nan_to_num(z_values, nan=0.0)
        if z_values.size == 0:
            return np.array([0.0, 1.0], dtype=float)
        if np.all(np.diff(z_values) >= 0):
            boundaries = np.concatenate(([0.0], z_values))
        else:
            boundaries = np.concatenate(([0.0], np.cumsum(np.abs(z_values))))
        return boundaries * 100.0
    if "layer_index" in df.columns:
        layer_index = df["layer_index"].to_numpy(dtype=float)
        layer_index = np.nan_to_num(layer_index, nan=0.0)
        if layer_index.size == 0:
            return np.array([0.0, 1.0], dtype=float)
        return np.arange(0.0, layer_index.size + 1, dtype=float) * 100.0

    z_values = np.arange(1, len(df) + 1, dtype=float)
    z_values = np.nan_to_num(z_values, nan=0.0)
    if z_values.size == 0:
        return np.array([0.0, 1.0], dtype=float)

    if np.any(np.diff(z_values) < 0):
        z_values = np.cumsum(np.abs(z_values))

    boundaries = np.concatenate(([0.0], np.cumsum(z_values)))
    return boundaries * 100.0


def _prepare_layer_values(values: np.ndarray, n_layers: int) -> np.ndarray:
    arr = values.astype(float, copy=False) if values.dtype.kind in {"i", "u", "f"} else values.copy()
    if arr.size > n_layers:
        arr = arr[:n_layers]
    elif arr.size < n_layers:
        padding = np.full(n_layers - arr.size, np.nan)
        arr = np.concatenate((padding, arr.astype(float, copy=False)))
    return arr