import re
from dataclasses import dataclass
from pathlib import Path
from typing import ClassVar

import matplotlib.pyplot as plt
import pandas as pd

from .base import SimStadtResults, detect_decimal
from .kpi import KPI
from .providers import Providers


@dataclass(repr=False)
class GreenWaterResults(SimStadtResults):
    """Results for GreenWater workflows (hourly time series with tree species breakdown)."""

    csv_identifier = "_greenwater"
    supported_providers: ClassVar[list[str]] = [Providers.GREEN_WATER]

    def _parse_results(self) -> pd.DataFrame:
        csv_decimal = detect_decimal(self.csv_path, 5, "tree area")
        df = pd.read_csv(
            self.csv_path,
            sep=";",
            header=[0, 1],
            skiprows=list(range(8)) + list(range(10, 13)),
            decimal=csv_decimal,
        )

        # Flatten multi-level columns: species names are forward-filled across sub-columns
        cols = df.columns.tolist()
        new_cols = []
        current_species = "All"
        species = []

        for col in cols:
            level0, level1 = col
            if pd.notna(level0) and not level0.startswith("Unnamed") and not level0.startswith("#"):
                current_species = level0
                species.append(current_species)
            new_cols.append(level1 if current_species == "All" else f"{current_species} - {level1}")

        df.columns = new_cols

        # Parse header metadata (tree counts, areas per species)
        with open(self.csv_path, encoding="utf-8") as csv:
            line = ""
            for _ in range(5):
                line = next(csv)
            tree_count = int(line.split(";")[1])
            line = next(csv)
            tree_area_header = re.split("[;,]", line.strip())
            total_tree_area = float(tree_area_header[1].replace(",", "."))
            if tree_count == 0:
                raise ValueError(f"CityGML {self.citygml} does not have any tree!")
            for _ in range(5):
                line = next(csv)
            species_count = [int(c) for c in re.split(";+", line.strip())[1:-1]]
            line = next(csv)
            species_area = [float(c.replace(",", ".")) for c in re.split(";+", line.strip())[1:-1]]

        species.pop(0)
        df.attrs = {
            "species": {
                name: {"count": count, "area": area}
                for name, count, area in zip(species, species_count, species_area)
            },
            "tree count": tree_count,
            "tree area": total_tree_area,
        }
        return df

    def df_with_time(self) -> pd.DataFrame:
        """Return the dataframe with a DatetimeIndex (year 2005, hourly)."""
        df = self.dataframe
        times = pd.date_range(start="2005-01-01 00:00", freq="1h", periods=8760)
        df = df.set_index(times)
        return df.drop(columns=["#"])

    @property
    def kpis(self) -> list[KPI]:
        df = self.dataframe
        rain = df["Rain [mm]"].sum()
        etc = df["Average ETc [mm]"].sum()
        cooling = df["Evaporative cooling [kWh/h]"].sum() / 1000
        irrigation_mm = df["Average irrigation [mm]"].sum()
        tree_area = df.attrs["tree area"]
        irrigation_m3 = irrigation_mm * tree_area / 1000
        return [
            KPI("Number of trees", df.attrs["tree count"], precision=0),
            KPI("Projected tree area", tree_area, "m²", precision=1),
            KPI("Rain", rain, "mm / a", precision=0),
            KPI("Average ETc", etc, "mm / a", precision=0),
            KPI("Irrigation", irrigation_m3, "m³ / a", precision=0),
            KPI("Evaporative cooling", cooling, "MWh / a", precision=0),
        ]

    @property
    def more_info(self) -> dict[str, list[KPI]]:
        df = self.dataframe
        all_info = {}
        for name, info in df.attrs["species"].items():
            count = info["count"]
            area = info["area"]
            irrigation_mm = df[f"{name} - Average irrigation [mm]"].sum()
            average_area = area / count
            all_info[f"_{name}_"] = [
                KPI("Number of trees", count, precision=0),
                KPI("Average tree area", average_area, "m²", precision=1),
                KPI("Average ETc", df[f"{name} - Average ETc [mm]"].sum(), "mm / a", precision=0),
                KPI("Irrigation", irrigation_mm, "mm / a", precision=0),
                KPI("Average irrigation", average_area * irrigation_mm / 1000, "m³ / a", precision=0),
            ]
        return all_info

    @property
    def diagrams(self) -> dict[str, Path]:
        rain_irrigation_png = self.workflow_path / "rain_and_irrigation.png"
        df_months = self.df_with_time()[["Rain [mm]", "Average irrigation [mm]"]].resample("ME").sum()
        fig, ax = plt.subplots()
        df_months = df_months.rename(
            columns={"Rain [mm]": "Niederschlag", "Average irrigation [mm]": "Künstliche Bewässerung"}
        )
        df_months.plot(kind="bar", stacked=True, ax=ax, ylabel="[mm]", figsize=(13, 6))
        ax.set_xticklabels([x.strftime("%b") for x in df_months.index], rotation=0)
        plt.savefig(rain_irrigation_png, bbox_inches="tight", dpi=300)
        plt.close(fig)
        return {"Rain and Irrigation": rain_irrigation_png}
