######################################################################
#             / ____|  ____|  __ \| |  | |  ____|  ____|             #
#            | |    | |__  | |__) | |__| | |__  | |__                #
#            | |    |  __| |  ___/|  __  |  __| |  __|               #
#            | |____| |____| |    | |  | | |____| |____              #
#             \_____|______|_|    |_|  |_|______|______|             #
######################################################################
#
# CEPHEE
# Copyright (C) 2024 Toulouse INP
#
# 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 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License for more details :
# <http://www.gnu.org/licenses/>.
#
######################################################################

# global
import os
import geopandas as gpd
import numpy as np
import rasterio
from shapely import Point

# local
from core.ModelCatchment import ModelCatchment
from core.Tools import is_raster, convert_geometry_to_multilinestring
from core.Data import polygonize_water_mask
from core.CrossSection import (import_XS_lines, optimize_XS_lines, remove_crossing_XS,
                           project_XS_lines, create_all_centerlines)

# optional
try:
    import h5py
    export_hdf_avail = True
except:
    export_hdf_avail = False
    print("no hdf5 librairy for binary export")


def load_project(project_path, params):
    """ reload existing catchment data stored as shp

    :param project_path: existing CEPHEE project path
    :type project_path: str, pickle, modelCatchment
    :param params: Parameters requires for computation
    :type params: Parameters
    """

    BV = ModelCatchment
    params.work_path = project_path
    params['I']['outlet_point'] = None
    BV.ordered_network = gpd.read_file(os.path.join(params.work_path, 'ordered_network.shp'))
    BV.hydro_network = BV.ordered_network
    BV.network_type = 'fromImport'
    params['I']['network_filepath'] = os.path.join(params.work_path, 'ordered_network.shp')

    BV.junction = gpd.read_file(os.path.join(params.work_path, 'junction.shp'))
    if os.path.isfile(os.path.join(project_path, 'XS_lines.shp')):
        params['I']['XS_filepath'] = os.path.join(project_path, 'XS_lines.shp')
    else:
        print('No cross sections available')

    lowest_point = None
    min_z = float('Inf')
    for geom in BV.ordered_network.geometry:
        for coord in geom.coords:
            if coord[2] < min_z:
                min_z = coord[2]
                lowest_point = coord

    params['I']['outlet_point'] = Point(lowest_point[0], lowest_point[1])
    BV.outlet_point = params['I']['outlet_point']
    BV.find_junctionAndOutlet(params['N']['minDistJunction'])  # trouve la position des confluences et les exutoires
    BV.renameReachFromJunction(params)

    if params['I']['riverbanks_filepath']:
        riverbanks_filepath = params['I']['riverbanks_filepath']
        if is_raster(params['I']['riverbanks_filepath']):
            riverbanks_filepath = polygonize_water_mask(riverbanks_filepath, params.work_path)
        bank_gdf = gpd.read_file(riverbanks_filepath)
        BV.outline = convert_geometry_to_multilinestring(bank_gdf)

    import_XS_lines(BV, params)
    if params['XS']['optimize_XS']:
        optimize_XS_lines(BV, method=params['XS']['optimization_method'],
                          max_iter=params['H']['MaxIter'], verbose=params.verbose)
        remove_crossing_XS(BV)
    project_XS_lines(BV, params)
    create_all_centerlines(BV, params.work_path)

    if os.path.isfile(os.path.join(project_path, 'flow_acc.tif')):
        acc_path = os.path.join(params.work_path, 'flow_acc.tif')

        acc_raster = rasterio.open(acc_path)
        acc = acc_raster.read(1)
        params['C']['resolution'] = acc_raster.transform[0]
        x_min, y_min, x_max, y_max = acc_raster.bounds
        X, Y = np.meshgrid(np.arange(x_min, x_max, params['C']['resolution']),
                           np.arange(y_min, y_max, params['C']['resolution']),
                           indexing='xy')
        Y = np.flipud(Y)
        # carte pour l'hydrologie
        BV.Map = {'X': X, 'Y': Y, 'fdir': None, 'acc': acc, 'mask': None, 'w1': None, 'd1': None, 'i2': None}
        BV.crs = acc_raster.crs

    return BV, params


def save_BV_results(BV,params,config_cal):


    # Exemple : définir ces variables avant
    # coord_pro, vh_pro, sim_name, wse, xy_h
    # raw_z, raw_V, raw_Q, raw_cs_xs, raw_xs, raw_cs
    # Unit = "SI Units" ou autre
    # Unsteady = True ou False
    if export_hdf_avail:
        output_file = params.work_path + config_cal +".hdf"
        with h5py.File(output_file, 'w') as f:

            root = f.create_group("Catchment")
            write_object_to_hdf5(root, BV)
        print(" Fichier HDF5 écrit avec succès ")

    return output_file

def write_object_to_hdf5(hdf_group, obj):
    #attrs_dict =vars(obj)
    new_dict = {attr: getattr(obj, attr) for attr in vars(obj)}
    for attr_name, attr_value in new_dict.items():
        if isinstance(attr_value, (int, float, str, bytes)):
            # Si c'est un attribut de base (int, float, string, bytes), on le stocke en attribut
            hdf_group.attrs[attr_name] = attr_value
        #elif isinstance(attr_value, tuple):
        elif isinstance(attr_value, list):
            # Si c'est une liste d'objets personnalisés
            if not attr_name in ('var','list_of_outlet'):
                if attr_name in ('section_list', 'reach'):
                    subgrp = hdf_group.create_group(attr_name)
                    for i, item in enumerate(attr_value):
                        item_grp = subgrp.create_group(f"{attr_name}_{i}")
                        if type(attr_value) == LineString or type(attr_value) == Point:
                            item =item.coords
                        write_object_to_hdf5(item_grp, item)
                elif attr_name in ('outline','elevationmin'):
                    subgrp = hdf_group.create_group(attr_name)
                    for i, item in enumerate(attr_value):
                        item_grp = subgrp.create_group(f"{attr_name}_{i}")
                        item = item.coords
                        write_object_to_hdf5(item_grp, item)

                elif attr_name in ('bank'):
                    hdf_group.create_dataset('distance', data=attr_value[0])
                    data = np.zeros((2,2))
                    data[0][0] = attr_value[1][0].x
                    data[0][1] = attr_value[1][0].y
                    data[1][0] = attr_value[1][1].x
                    data[1][1] = attr_value[1][1].y
                    hdf_group.create_dataset('bank_point', data=data)

                elif attr_name in ('normal'):
                    data = np.array([attr_value[0],attr_value[1]])
                    hdf_group.create_dataset('normal', data=data)
                else:
                    attr_value = np.array(attr_value)
                    hdf_group.create_dataset(attr_name, data=attr_value)
            else:
                #attr_value = np.array(attr_value)
                #print(attr_value)
                hdf_group.create_dataset(attr_name, data=str(attr_value))
        elif isinstance(attr_value, np.ndarray):
            # Si c'est un tableau numpy, on le stocke comme dataset
            hdf_group.create_dataset(attr_name, data=attr_value)
        elif type(attr_value) == LineString or type(attr_value) == Point:
            hdf_group.attrs[attr_name] = attr_value.coords
        elif isinstance(attr_value, Coord):
            hdf_group.attrs['array'] = attr_value.array
            hdf_group.attrs['values'] = attr_value.values




def read_object_from_hdf5(hdf_group, obj_class, subgroups_types=None):
    """
    Reconstructs an object from an HDF5 group, handling lists of objects recursively.

    :param hdf_group: HDF5 group containing the data for the object
    :param obj_class: The class to instantiate (e.g., BV)
    :param subgroups_types: A dictionary mapping subgroup names to class types
    :return: An instance of the given class populated with data from the HDF5 group
    """
    # Initialize an empty dictionary for object attributes
    obj_data = {}

    # Read attributes from the HDF5 group
    for attr_name, attr_value in hdf_group.attrs.items():
        obj_data[attr_name] = attr_value

    # Read subgroups (for lists or nested objects)
    for subgroup_name, subgroup in hdf_group.items():
        if isinstance(subgroup, h5py.Group):
            # If this is a list (e.g., reach_0, reach_1, etc.)
            if subgroup_name.endswith("_0"):  # This assumes lists are named like reach_0, reach_1, ...
                sublist = []
                for sub_item_name in subgroup:
                    sub_item = subgroup[sub_item_name]
                    if subgroup_name[:4] =='cont':
                        sub_obj = sub_item
                    else:
                        sub_obj = read_object_from_hdf5(sub_item, obj_class, subgroups_types)  # Recursive call
                    sublist.append(sub_obj)
                obj_data[subgroup_name] = sublist
            else:
                # If it's a nested object, determine the class to use for reconstruction
                # Look up the class type for this subgroup in the dictionary

                subgroup_class = subgroups_types.get(subgroup_name, obj_class)
                nested_obj = read_object_from_hdf5(subgroup, subgroup_class, subgroups_types)
                obj_data[subgroup_name] = nested_obj



        #print(subgroup_name)

    # Create the object from the dictionary of attributes

    obj = obj_class(**obj_data)

    return obj
