import h5py
import pandas as pd
import numpy as np
import re

import matplotlib as plt

from Model import Model
from fitter import Fitter
from Data import Data

file = "ltcc_current.h5"

dfs_by_sex_tag_spid = {}


def print_attrs(name, obj):
    #  print(f"\nAttributes for {name}:")
    #  for key, val in obj.attrs.items():
    #  print(f"  {key}: {val}")
    pass


with h5py.File(file, "r") as h5_file:
    for eid in h5_file.keys():
        attributes = h5_file[eid].attrs
        sex = attributes.get("sex")
        tag = attributes.get("tag")
        spid = attributes.get("spid")

        key = f"{sex}_{tag}_{spid}"

        if key not in dfs_by_sex_tag_spid:
            dfs_by_sex_tag_spid[key] = pd.DataFrame()

        row_data = {"experiment_id": eid, "sex": sex, "tag": tag, "spid": spid}
        temp_df = pd.DataFrame([row_data])

        dfs_by_sex_tag_spid[key] = pd.concat(
            [dfs_by_sex_tag_spid[key], temp_df], ignore_index=True
        )

#  for key, df in dfs_by_sex_tag_spid.items():
#  print(f"DataFrame for {key}:")
#  print(df)
#  print()


def fit_and_plot_dataframes(dfs_by_sex_tag_spid):
    for key, df in dfs_by_sex_tag_spid.items():
        print(f"Fitting and plotting data for {key}...")

        combined_current = []
        combined_time = []

        for eid in df["experiment_id"].tolist():
            data = Data(file, group_key=eid)
            combined_current.append(data.current)
            combined_time.append(data.current_t)

        combined_current = np.concatenate(combined_current)
        combined_time = np.concatenate(combined_time)

        # Sort by time for consistency - inspired by data
        sorted_indices = np.argsort(combined_time)
        combined_time = combined_time[sorted_indices]
        combined_current = combined_current[sorted_indices]

        combined_data = Data(file, group_key=eid)
        combined_data.current = combined_current
        combined_data.current_t = combined_time

        fit = Fitter(Model, combined_data)
        fit.optimize()
        res, fig = fit.optimize()

        plt.figure()
        plt.title(f"Fit results for {key}")

        for i, eid in enumerate(df["experiment_id"].tolist()):
            plt.plot(combined_time, combined_current,
                     label=f"Experiment {eid}")

        plt.plot(
            combined_data.current_t, combined_data.current, "k-",
            label="Combined Fit"
        )

        plt.legend()
        plt.xlabel("Time")
        plt.ylabel("Current")

        key_cleaned = re.sub(r"[^\w.-]", "", key)
        plt.savefig(f"combined_plot_{key_cleaned}.png")
        plt.savefig(f"combined_plot_{key_cleaned}.pdf")
        plt.close()

        fit_hist = pd.DataFrame.from_dict(fit.fit_results, orient="index").T
        fit_hist.index.name = "Iterations"

        res_filename = f"combined_fit_results_{key}.csv"
        res_filename = res_filename.replace(" ", "_").replace(":", "-")
        fit_hist.to_csv(res_filename, index=True)

        print(f"Finished fitting for {key}.")


fit_and_plot_dataframes(dfs_by_sex_tag_spid)

"""
def fit_data():
    filename = "ltcc_current.h5"
    with h5py.File(filename, "r") as h5:
        eids = list(h5.keys())

    for eid in eids:
        data = Data(filename, group_key=eid)
        fit = Fitter(Model, data)
        fit.optimize()
        res, fig = fit.optimize()

        fit_hist = pd.DataFrame.from_dict(fit.fit_results, orient="index").T
        fit_hist.index.name = "Iterations"

        res_filename = f"fit_results_{eid}.csv"
        res_filename = res_filename.replace(" ", "_").replace(":", "-")
        fit_hist.to_csv(res_filename, index=True)

        eid_cleaned = re.sub(r"[^w.-]", "", eid)  # Eemaldab kõik eritähed
        fig.savefig(f"plot_{eid_cleaned}.png")
        fig.savefig(f"plot_{eid_cleaned}.pdf")


fit_data()

"""