import pathlib
import zipfile

import matplotlib.pyplot as plt
import pandas as pd
from tabulate import tabulate


def decompress_zipfile(zipfile_path, destination_folder="data/out"):
    """Decompress a file into the desired folder."""    
    with zipfile.ZipFile(zipfile_path, 'r') as zip_file:
        for compressed_file in zip_file.namelist():
            zip_file.extract(compressed_file, destination_folder)
            print(f"Decompressed file {compressed_file} from {zip_file.filename}")
            

def extract_dataframes_from_iso(file_path):
    df_list = []  # List to store the extracted DataFrames
    header_rows = []  # List to store the row indices of the headers
    
    # Read the ISO file into a DataFrame using fixed-width format
    df = pd.read_fwf(file_path, header=None)
    
    # Loop through each row and find the header rows
    for index, row in df.iterrows():
        # Check if all the values in the row are non-empty (not NaN)
        if row.notnull().all():
            header_rows.append(index)
    
    # Add the last row index as the end of the last header section
    header_rows.append(len(df))
    
    # HACK: skip empty header rows
    header_rows = header_rows[1::2]
    
    # Split the DataFrame based on the header rows
    for i in range(len(header_rows) - 1):
        start = header_rows[i]
        end = header_rows[i + 1]
        # Extract the DataFrame for this section
        df_section = df.iloc[start:end, :]
        # Set the first row as the header for the DataFrame
        df_section.columns = df_section.iloc[0]
        # Drop the first row, as it's now the header
        df_section = df_section.iloc[1:]
        
        # HACK: clean malformatted column data
        malformed_col_name = "EEP          log10_isochrone_age_yr                    initial_mass                       star_mass"
        malformed_data = df_section[malformed_col_name].str.split(expand=True)
        epp, age, initial_mass, star_mass = [malformed_data[i] for i in range(4)]
        df_section["EPP"] = epp
        df_section["log10_isochrone_age_yr"] = age
        df_section["initial_mass"] = initial_mass
        df_section["star_mass"] = star_mass
        df_section.drop(columns=[malformed_col_name, "#"], inplace=True)
        df_section.drop(df_section.tail(2).index, inplace=True)
        
        # Append the DataFrame to the list
        df_list.append(df_section.astype(float))
    
    return df_list


def find_min_max_values(dataframes, columns_to_check):
    """Find the minimum and maximum values for specified columns across all DataFrames."""
    
    result = {}
    
    for column in columns_to_check:
        all_column_values = []
        for df in dataframes:
            if column in df.columns:
                all_column_values.extend(df[column].astype(float).values)
        
        if all_column_values:
            min_value = min(all_column_values)
            max_value = max(all_column_values)
            result[column] = (min_value, max_value)
            
    rows = [{'Data': col, 'Min': val[0], 'Max': val[1]} for col, val in result.items()]
    return pd.DataFrame(rows)


def filter_data_within_Teff_and_g_range(df, min_log_Teff, max_log_Teff, min_log_g, max_log_g):
    # Convert the "log_Teff" and "log_g" columns to float type
    df["log_Teff"] = df["log_Teff"].astype(float)
    df["log_g"] = df["log_g"].astype(float)

    # Create the boolean mask for the condition (data within the range)
    mask = (
        (df["log_Teff"] >= min_log_Teff)
        & (df["log_Teff"] <= max_log_Teff)
        & (df["log_g"] >= min_log_g)
        & (df["log_g"] <= max_log_g)
    )
    
    # Apply the mask to the DataFrame and return the filtered data
    filtered_data = df[mask]
    
    return filtered_data


def compute_mean_core_hidrogen_quantity_within_Teff_and_g_range(
    df, min_log_Teff, max_log_Teff, min_log_g, max_log_g
):
        # Convert the "log_Teff" and "log_g" columns to float type
    df["log_Teff"] = df["log_Teff"].astype(float)
    df["log_g"] = df["log_g"].astype(float)

    # Create the boolean mask for the condition (data within the range)
    mask = (
        (df["log_Teff"] >= min_log_Teff)
        & (df["log_Teff"] <= max_log_Teff)
        & (df["log_g"] >= min_log_g)
        & (df["log_g"] <= max_log_g)
    )
    
    filtered_data = df[mask]
    return filtered_data["center_h1"].mean()
    


def plot_isochrones_with_limits(df_list, limits, zoom_pos, zoom_limits, title="H-R diagram", ax=None):
    if ax is None:
        _, ax = plt.subplots()
    zoom_ax = ax.inset_axes(zoom_pos)
        
    for data in df_list:
        observed_data = filter_data_within_Teff_and_g_range(data, *limits)

        ax.plot(data["log_Teff"], data["log_L"], "r--")
        ax.plot(observed_data["log_Teff"], observed_data["log_L"], "b*")
        
        zoom_ax.plot(data["log_Teff"], data["log_L"], "r--")
        zoom_ax.plot(observed_data["log_Teff"], observed_data["log_L"], "b*")

    ax.set_title(title)
    ax.set_xlabel("log_Teff")
    ax.set_ylabel("log_L")
    
    zoom_ax.set_xlim(*zoom_limits[0:2])
    zoom_ax.set_ylim(*zoom_limits[2:])
    ax.indicate_inset_zoom(zoom_ax, edgecolor="black")
    
    plt.gca().invert_xaxis()
    
    return ax


def read_trackfile_as_dataframe(trackfile, skiprows=11, delim_whitespace=True):
    """Read a ``.track.epp`` file and return a ``DataFrame`` representing it."""
    return (
        pd.read_csv(trackfile, skiprows=skiprows, delim_whitespace=delim_whitespace)
        .shift(periods=1, axis=1)
        .drop(columns="#")
    )


def read_isochrone_as_dataframe(isochronefile, skiprows=12, delim_whitespace=True):
    """Read a ``.track.epp`` file and return a ``DataFrame`` representing it."""
    return (
        pd.read_csv(isochronefile, skiprows=skiprows, delim_whitespace=delim_whitespace)
        #.shift(periods=1, axis=1)
        #.drop(columns="#")
    )

for zipfile_path in pathlib.Path().glob("data/**/*.zip"):
    decompress_zipfile(zipfile_path)

import numpy as np


# Temperature limits
base_Teff = 8170
delta_Teff = 100
min_log_Teff = np.log10(base_Teff - delta_Teff)
max_log_Teff = np.log10(base_Teff + delta_Teff)

# Gravity limits
base_log_g = 3.90
delta_log_g = 0.30
min_log_g = base_log_g - delta_log_g
max_log_g = base_log_g + delta_log_g

solar_metallicity_data = extract_dataframes_from_iso("data/out/isochrone_fe_over_h_0.iso")
solar_metallicity_stats = find_min_max_values(
    solar_metallicity_data, ["star_mass", "log_L", "log_Teff", "log10_isochrone_age_yr"]
)
print(solar_metallicity_stats)
print("=== Tabla 1 - Resumen de resultados para una estrella con metalicidad solar ===")

min_Teff, max_Teff = 10 ** min_log_Teff, 10 ** max_log_Teff
min_L, max_L = 10 ** 1.1, 10 ** 1.9

print(f"Rango de temperatura efectiva: [{min_Teff:.2f}, {max_Teff:.2f}] Kelvin")
print(f"Rango de luminosidad efectiva: [{min_L:.2f}, {max_L:.2f}] unidades solares")

mean_quantity_center_h1_list = []

for isochrone in solar_metallicity_data:
    mean_h1 = compute_mean_core_hidrogen_quantity_within_Teff_and_g_range(
        isochrone, min_log_Teff, max_log_Teff, min_log_g, max_log_g
    )
    if mean_h1 is not np.nan:
        mean_quantity_center_h1_list.append(mean_h1)
    
mean_quantity_center_h1 = np.mean(mean_quantity_center_h1_list)
print(f"Cantidad media de hidrógeno en el nucleo {mean_quantity_center_h1*100:.2f}%")

ax = plot_isochrones_with_limits(
    solar_metallicity_data[::2], limits=[min_log_Teff, max_log_Teff, min_log_g, max_log_g],
    zoom_pos=[-0.6, 0.82, 0.5, 0.5],
    zoom_limits=[3.93, 3.89, 1.0, 2.0],
    title="Figura 1 - Diagrama HR para estrella con metalicidad solar"
)
plt.show()

subsolar_metallicity_data = extract_dataframes_from_iso("data/out/isochrone_fe_over_h_minus_1.iso")
subsolar_metallicity_stats = find_min_max_values(
    subsolar_metallicity_data, ["star_mass", "log_L", "log_Teff", "log10_isochrone_age_yr"]
)
print(subsolar_metallicity_stats)
print("=== Tabla 2 - Resumen de resultados para una estrella con metalicidad subsolar ===")

mean_quantity_center_h1_list = []

for isochrone in subsolar_metallicity_data:
    mean_h1 = compute_mean_core_hidrogen_quantity_within_Teff_and_g_range(
        isochrone, min_log_Teff, max_log_Teff, min_log_g, max_log_g
    )
    if mean_h1 is not np.nan:
        mean_quantity_center_h1_list.append(mean_h1)
    
mean_quantity_center_h1 = np.mean(mean_quantity_center_h1_list)
print(f"Cantidad media de hidrógeno en el nucleo {mean_quantity_center_h1*100:.2f}%")

ax = plot_isochrones_with_limits(
    subsolar_metallicity_data[::2], limits=[min_log_Teff, max_log_Teff, min_log_g, max_log_g],
    zoom_pos=[-0.6, 0.82, 0.5, 0.5],
    zoom_limits=[3.88, 3.94, 0.8, 1.85],
    title="Figura 2 - Diagrama HR para estrella con metalicidad subsolar"
)
plt.show()

import pandas as pd

evolutionary_track_M2 = read_trackfile_as_dataframe("./data/out/0020000M.track.eep")
evolutionary_track_M15 = read_trackfile_as_dataframe("./data/out/0150000M.track.eep")

import matplotlib.pyplot as plt

fig, ax = plt.subplots()
ax.plot(evolutionary_track_M2["log_Teff"], evolutionary_track_M2["log_L"], color="blue", label=r'2$M_{\odot}$')
ax.plot(evolutionary_track_M15["log_Teff"], evolutionary_track_M15["log_L"], color="red", label=r'15$M_{\odot}$')
ax.set_title("Figura 3 - Trazas evolutivas")
ax.set_xlabel("log_Teff")
ax.set_ylabel("log_L")
ax.legend()
plt.gca().invert_xaxis()