"""
kansim.py - Minimal Monte Carlo simulation framework
"""

import os
import re
import numpy as np
import pandas as pd
import plotly.express as px
from datetime import datetime, timezone
from scipy import stats as scipy_stats

__all__ = [
    "Simulation",
    "SimulationResult",
    "Normal", "LogNormal", "Uniform", "Triangular", "BetaPERT",
    "Beta", "GeneralizedBeta", "Exponential", "Gamma", "Weibull",
    "Pareto", "Poisson", "Binomial", "NegativeBinomial", "StudentT",
    "PearsonIII", "ExtremeValue", "Discrete", "Cumulative",
    "SampledResults", "Boolean",
]


# ---------------------------------------------------------------------------
# Distribution functions — each returns a single sampled float
# ---------------------------------------------------------------------------

def Normal(mean=0, std=1, min=None, max=None):
    v = np.random.normal(mean, std)
    return float(np.clip(v, min, max) if (min is not None or max is not None) else v)

def LogNormal(mean=1, std=0.5, min=None, max=None):
    """True (arithmetic) mean and std."""
    sigma2 = np.log(1 + (std / mean) ** 2)
    mu = np.log(mean) - sigma2 / 2
    v = np.random.lognormal(mu, np.sqrt(sigma2))
    return float(np.clip(v, min, max) if (min is not None or max is not None) else v)

def Uniform(min=0, max=1, log=False):
    if log:
        return float(np.exp(np.random.uniform(np.log(min), np.log(max))))
    return float(np.random.uniform(min, max))

def Triangular(min=0, most_likely=0.5, max=1):
    return float(np.random.triangular(min, most_likely, max))

def BetaPERT(min=0, most_likely=0.5, max=1):
    """PERT-weighted beta distribution."""
    mean = (min + 4 * most_likely + max) / 6
    if max == min:
        return float(most_likely)
    alpha = 6 * (mean - min) / (max - min)
    beta_ = 6 * (max - mean) / (max - min)
    return float(min + (max - min) * np.random.beta(alpha, beta_))

def Beta(successes=2, failures=2):
    """Beta distribution parameterized by successes/failures (0–1 range)."""
    return float(np.random.beta(successes, failures))

def GeneralizedBeta(mean, std, min=0, max=1):
    """Beta scaled to [min, max] via mean and std."""
    r = max - min
    mu = (mean - min) / r
    alpha = mu * (mu * (1 - mu) / (std / r) ** 2 - 1)
    beta_ = (1 - mu) * (mu * (1 - mu) / (std / r) ** 2 - 1)
    return float(min + r * np.random.beta(alpha, beta_))

def Exponential(mean=1):
    return float(np.random.exponential(mean))

def Gamma(mean=1, std=1, min=None, max=None):
    shape = (mean / std) ** 2
    scale = std ** 2 / mean
    v = np.random.gamma(shape, scale)
    return float(np.clip(v, min, max) if (min is not None or max is not None) else v)

def Weibull(scale=1, shape=1, min=0, max=None):
    """scale = characteristic life (mean - min), shape = slope (k)."""
    v = min + scale * np.random.weibull(shape)
    return float(np.clip(v, None, max) if max is not None else v)

def Pareto(shape=1, mode=1, max=None):
    v = mode * (np.random.pareto(shape) + 1)
    return float(np.clip(v, None, max) if max is not None else v)

def Poisson(expected=1):
    return float(np.random.poisson(expected))

def Binomial(n=10, p=0.5):
    return float(np.random.binomial(n, p))

def NegativeBinomial(successes=5, p=0.5):
    return float(np.random.negative_binomial(successes, p))

def StudentT(df=10):
    return float(np.random.standard_t(df))

def PearsonIII(location=0, scale=1, shape=1):
    """Pearson Type III (shifted gamma)."""
    return float(scipy_stats.pearson3.rvs(shape, loc=location, scale=scale))

def ExtremeValue(location=0, scale=1):
    """Gumbel (Type I extreme value) distribution."""
    return float(np.random.gumbel(location, scale))

def Discrete(values, probs):
    """Discrete distribution. values and probs are lists."""
    probs = np.array(probs, dtype=float)
    probs /= probs.sum()
    return float(np.random.choice(values, p=probs))

def Cumulative(values, probs, log=False):
    """Empirical CDF via inverse transform. probs are cumulative (0..1)."""
    u = np.random.uniform()
    v = float(np.interp(u, probs, values))
    return np.exp(v) if log else v

def SampledResults(values):
    """Non-parametric: resample from observed data."""
    return float(np.random.choice(values))

def Boolean(p=0.5):
    return bool(np.random.random() < p)


# ---------------------------------------------------------------------------
# SimulationResult
# ---------------------------------------------------------------------------

def _make_result_id():
    """Human readable, URL-safe timestamp-based ID e.g. 2024-03-10-143022"""
    return datetime.now(timezone.utc).strftime("%Y-%m-%d-%H%M%S")

def _safe_path(result_id, filename):
    """Resolve a sandboxed path under results/<result_id>/, block path traversal."""
    base = os.path.abspath(os.path.join("results", result_id))
    target = os.path.abspath(os.path.join(base, filename))
    if not target.startswith(base + os.sep) and target != base:
        raise ValueError(f"Unsafe path detected: {filename}")
    return target

def _sanitize_filename(name):
    """Strip any path separators or unsafe characters from a filename."""
    return re.sub(r"[^\w\-.]", "_", os.path.basename(name))


class SimulationResult:
    """
    Returned by Simulation.run(). Holds the DataFrame and result_id.
    All saves write to results/<result_id>/.
    """

    def __init__(self, df, result_id):
        self.df = df
        self.result_id = result_id
        self._dir = os.path.join("results", result_id)
        os.makedirs(self._dir, exist_ok=True)

    def _resolve(self, filename):
        return _safe_path(self.result_id, _sanitize_filename(filename))

    def _show_or_save(self, fig, save):
        if save:
            path = self._resolve(save)
            fig.write_image(path)
        else:
            fig.show()

    def save_csv(self, filename="data.csv"):
        path = self._resolve(filename)
        self.df.to_csv(path, index=False)
        return path

    def save_parquet(self, filename="data.parquet"):
        path = self._resolve(filename)
        self.df.to_parquet(path, index=False)
        return path

    def plot_histogram(self, col, df=None, bins=50, title=None, save=None):
        d = df if df is not None else self.df
        fig = px.histogram(d, x=col, nbins=bins, title=title or f"Distribution of {col}")
        self._show_or_save(fig, save)

    def plot_cdf(self, col, df=None, title=None, save=None):
        d = df if df is not None else self.df
        data = d[col].dropna().sort_values()
        cdf = np.arange(1, len(data) + 1) / len(data)
        fig = px.line(x=data, y=cdf, labels={"x": col, "y": "Cumulative Probability"},
                      title=title or f"CDF of {col}")
        self._show_or_save(fig, save)

    def plot_scatter(self, x_col, y_col, df=None, title=None, opacity=0.3, save=None):
        d = df if df is not None else self.df
        fig = px.scatter(d, x=x_col, y=y_col, opacity=opacity,
                         title=title or f"{y_col} vs {x_col}")
        self._show_or_save(fig, save)

    def plot_tornado(self, target, df=None, title=None, save=None):
        """Spearman rank correlation of all numeric columns vs target."""
        d = df if df is not None else self.df
        numeric = d.select_dtypes(include="number").drop(columns=[target], errors="ignore")
        correlations = {
            col: scipy_stats.spearmanr(d[col].dropna(), d[target].dropna())[0]
            for col in numeric.columns
        }
        corr_df = pd.DataFrame(
            sorted(correlations.items(), key=lambda x: x[1]),
            columns=["variable", "correlation"]
        )
        fig = px.bar(corr_df, x="correlation", y="variable", orientation="h",
                     color="correlation", color_continuous_scale="RdYlGn",
                     title=title or f"Sensitivity Tornado — {target}")
        fig.add_vline(x=0, line_color="black", line_width=1)
        self._show_or_save(fig, save)

    def __repr__(self):
        return f"SimulationResult(result_id={self.result_id!r}, runs={len(self.df)}, columns={list(self.df.columns)})"


# ---------------------------------------------------------------------------
# Base class
# ---------------------------------------------------------------------------

class Simulation:
    """
    Subclass this. Define all variables (sampled + computed) in __init__.
    Call MyModel.run(n=10_000) to get a SimulationResult.

    from kansim import Simulation, Normal, Beta, Triangular
    """

    @classmethod
    def run(cls, n=1000, seed=None):
        if seed is not None:
            np.random.seed(seed)
        records = []
        for _ in range(n):
            instance = cls()
            records.append(instance.__dict__)
        df = pd.DataFrame(records)
        return SimulationResult(df, _make_result_id())

