"""
init-mater-project

Copyright (C) 2025 [Lauranne Sarribouette] <lauranne.sarribouette@univ-grenoble-alpes.fr>

SPDX-License-Identifier: LGPL-3.0-or-later
"""

from pathlib import Path
from typing import List, Dict, Tuple, Optional
from enum import Enum
import pandas as pd
import streamlit as st


class OutputStructure(Enum):
    """Output folder structure modes"""

    SINGLE_SIMULATION = "single"
    MULTIPLE_SIMULATIONS = "multiple"
    INVALID = "invalid"


@st.cache_data(ttl=600)
def detect_output_structure(
    output_folder: Path,
) -> Tuple[OutputStructure, Optional[str]]:
    """Detect output folder structure automatically"""
    if not output_folder.exists():
        return OutputStructure.INVALID, f"Output folder does not exist: {output_folder}"

    if not output_folder.is_dir():
        return (
            OutputStructure.INVALID,
            f"Output path is not a directory: {output_folder}",
        )

    subdirs = [
        d for d in output_folder.iterdir() if d.is_dir() and not d.name.startswith(".")
    ]

    if not subdirs:
        return OutputStructure.INVALID, f"No subdirectories found in {output_folder}"

    # Check depth 1: do subdirectories contain .mater files directly?
    mater_files_depth_1 = []
    for subdir in subdirs:
        mater_files = list(subdir.glob("*.mater"))
        if mater_files:
            mater_files_depth_1.append(subdir)

    if mater_files_depth_1:
        return OutputStructure.SINGLE_SIMULATION, None

    # Check depth 2: do subdirectories contain other subdirectories with .mater files?
    simulations_found = []
    for subdir in subdirs:
        sub_subdirs = [
            d for d in subdir.iterdir() if d.is_dir() and not d.name.startswith(".")
        ]

        for sub_subdir in sub_subdirs:
            mater_files = list(sub_subdir.glob("*.mater"))
            if mater_files:
                simulations_found.append(subdir)
                break

    if simulations_found:
        return OutputStructure.MULTIPLE_SIMULATIONS, None

    error_msg = f"""Invalid output folder structure.

Expected structures:

Option 1 (single simulation):
{output_folder.name}/
  ├── variable1/
  │   └── *.mater
  └── variable2/
      └── *.mater

Option 2 (multiple simulations):
{output_folder.name}/
  ├── simulation_1/
  │   ├── variable1/*.mater
  │   └── variable2/*.mater
  └── simulation_2/
      ├── variable1/*.mater
      └── variable2/*.mater

No .mater files found. Run a simulation first:
  uv run mater-cli simulation run --example
"""

    return OutputStructure.INVALID, error_msg


def get_available_simulations(output_folder: Path) -> List[str]:
    """Get list of available simulation folders"""
    if not output_folder.exists():
        return []

    return [
        d.name
        for d in output_folder.iterdir()
        if d.is_dir() and not d.name.startswith(".")
    ]


def get_available_variables(simulation_path: Path) -> List[str]:
    """Get list of available variables in a simulation"""
    if not simulation_path.exists():
        return []

    return [
        d.name
        for d in simulation_path.iterdir()
        if d.is_dir() and not d.name.startswith(".")
    ]


@st.cache_data(ttl=600)
def load_variable_data(simulation_path: Path, variable_name: str) -> pd.DataFrame:
    """
    Load variable data from simulation and aggregate age cohorts if present

    Args:
        simulation_path: Path to simulation folder
        variable_name: Name of the variable to load

    Returns:
        DataFrame with variable data (age_cohort aggregated if present)
    """
    variable_path = simulation_path / variable_name
    df = _load_mater_files(variable_path, variable_name)

    if df is None or df.empty:
        return pd.DataFrame()

    if "age_cohort" in df.index.names:
        level_list = [name for name in df.index.names if name != "age_cohort"]
        df = df.groupby(level=level_list).sum()

    return df


def _load_mater_files(variable_path: Path, variable_name: str) -> pd.DataFrame:
    """Load variable data from .mater files"""
    if not variable_path.exists():
        return pd.DataFrame()

    mater_files = list(variable_path.glob("*.mater"))

    if not mater_files:
        return pd.DataFrame()

    single_file = variable_path / f"{variable_name}.mater"
    if single_file.exists() and len(mater_files) == 1:
        return pd.read_parquet(single_file)

    dfs = [pd.read_parquet(file) for file in sorted(mater_files)]

    if not dfs:
        return pd.DataFrame()

    return pd.concat(dfs, ignore_index=False).sort_index()


def filter_data(
    df: pd.DataFrame,
    time_range: Optional[Tuple[pd.Timestamp, pd.Timestamp]],
    dimension_filters: Dict[str, List[str]],
    aggregate_all: bool = False,
) -> pd.DataFrame:
    """Filter DataFrame by time range and dimension values"""
    if df.empty:
        return df

    filtered = df.copy()

    if time_range and time_range[0] and time_range[1]:
        if "time" in filtered.index.names:
            time_level = filtered.index.get_level_values("time")
            mask = (time_level >= time_range[0]) & (time_level <= time_range[1])
            filtered = filtered[mask]

    for dim_name, values in dimension_filters.items():
        if dim_name not in filtered.index.names or not values:
            continue

        filtered = filtered[filtered.index.get_level_values(dim_name).isin(values)]

    return filtered


def count_series(df: pd.DataFrame) -> int:
    """Count number of unique series (combinations) in DataFrame"""
    if df.empty:
        return 0

    dims = [name for name in df.index.names if name not in ["time"]]

    if not dims:
        return 1

    return df.index.to_frame()[dims].drop_duplicates().shape[0]


def prepare_treemap_data(df: pd.DataFrame, primary_dim: str) -> pd.DataFrame:
    """Prepare data for single-dimension treemap"""
    if df.empty:
        return pd.DataFrame()

    if isinstance(df, pd.Series):
        df = df.to_frame(name="value")

    df_reset = df.reset_index()

    if "value" not in df_reset.columns:
        value_cols = [col for col in df_reset.columns if col not in [primary_dim]]
        if value_cols:
            df_reset = df_reset.rename(columns={value_cols[0]: "value"})

    df_clean = df_reset[[primary_dim, "value"]].copy()

    grouped = df_clean.groupby([primary_dim])["value"].sum().reset_index()

    treemap_data = [
        {
            "labels": str(row[primary_dim]),
            "parents": "",
            "values": row["value"],
            "ids": str(row[primary_dim]),
        }
        for _, row in grouped.iterrows()
    ]

    return pd.DataFrame(treemap_data)


def validate_dimension_filters(
    dimension_filters: Dict[str, List[str]],
) -> Tuple[bool, Optional[str]]:
    """Validate that each dimension has at least one value selected"""
    for dim_name, values in dimension_filters.items():
        if not values or len(values) == 0:
            return False, f"Please select at least one value for {dim_name}"

    return True, None
