"""
QF-Lib Quant Analytics — computes full strategy performance report.

Produces JSON with:
  - equityCurve: daily equity from trade history
  - monthlyReturns: heatmap matrix (years × months)
  - yearlyReturns: bar chart data with mean
  - monthlyReturnDistribution: histogram bins
  - qqPlot: theoretical vs observed quantiles
  - rollingStats: 6-month rolling return + volatility
"""

import json, math
from datetime import datetime, timedelta
from collections import defaultdict, OrderedDict
from typing import Optional

MONTHS = ["Jan", "Feb", "Mar", "Apr", "May", "Jun",
          "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]

def compute_daily_equity(trades: list[dict], start_equity: float = 100.0) -> list[dict]:
    """Build daily equity curve from trade PnL history."""
    daily = defaultdict(float)
    for t in trades:
        try:
            ts = t.get("time", "")
            if "T" in ts:
                date = ts[:10]
            elif " " in ts:
                date = ts.split(" ")[0]
            elif len(ts) >= 10:
                date = ts[:10]
            else:
                continue
            pnl = float(t.get("pnl", 0))
            daily[date] += pnl
        except (ValueError, KeyError):
            continue
    
    dates = sorted(daily.keys())
    if not dates:
        return [{"date": "2024-01-01", "value": start_equity}]
    
    equity = start_equity
    curve = []
    # Fill from first trade date to last
    first = datetime.strptime(dates[0], "%Y-%m-%d")
    last = datetime.strptime(dates[-1], "%Y-%m-%d")
    current = first
    while current <= last:
        d = current.strftime("%Y-%m-%d")
        if d in daily:
            equity += daily[d]
        curve.append({"date": d, "value": round(equity, 4)})
        current += timedelta(days=1)
    return curve

def compute_monthly_returns(equity_curve: list[dict]) -> dict:
    """Compute monthly returns from daily equity curve."""
    if len(equity_curve) < 2:
        return {"years": [], "months": MONTHS, "matrix": []}
    
    # Group by year-month
    monthly = OrderedDict()
    for pt in equity_curve:
        d = datetime.strptime(pt["date"], "%Y-%m-%d")
        ym = f"{d.year}-{d.month:02d}"
        if ym not in monthly:
            monthly[ym] = {"first": pt["value"], "last": pt["value"], "date": pt["date"]}
        monthly[ym]["last"] = pt["value"]
        monthly[ym]["date"] = pt["date"]
    
    # Compute returns
    months_data = []
    prev_value = None
    for ym, data in monthly.items():
        if prev_value is not None and prev_value > 0:
            ret = ((data["last"] / prev_value) - 1) * 100
        else:
            ret = None
        prev_value = data["last"]
        year = int(ym[:4])
        month = int(ym[5:7])
        months_data.append({"year": year, "month": month, "return": ret})
    
    if not months_data:
        return {"years": [], "months": MONTHS, "matrix": []}
    
    years = sorted(set(m["year"] for m in months_data), reverse=True)
    matrix = []
    for yr in years:
        row = [None] * 12
        for m in months_data:
            if m["year"] == yr:
                v = m["return"]
                row[m["month"] - 1] = round(v, 1) if v is not None else None
        matrix.append(row)
    
    return {"years": years, "months": MONTHS, "matrix": matrix}

def compute_yearly_returns(monthly_data: dict) -> tuple[list[dict], float]:
    """Compute yearly returns from monthly returns matrix."""
    years = monthly_data.get("years", [])
    matrix = monthly_data.get("matrix", [])
    yearly = []
    
    for i, yr in enumerate(years):
        total = 1.0
        row = matrix[i]
        has_data = False
        for v in row:
            if v is not None:
                total *= (1 + v / 100)
                has_data = True
        if has_data:
            ret = round((total - 1) * 100, 1)
            yearly.append({"year": yr, "return": ret})
    
    if not yearly:
        return [], 0.0
    
    mean = round(sum(r["return"] for r in yearly) / len(yearly), 1)
    return yearly, mean

def compute_return_distribution(monthly_data: dict) -> dict:
    """Compute histogram of monthly returns for distribution chart."""
    matrix = monthly_data.get("matrix", [])
    all_returns = []
    for row in matrix:
        for v in row:
            if v is not None:
                all_returns.append(v)
    
    if not all_returns:
        return {"bins": [], "mean": 0.0}
    
    mean = round(sum(all_returns) / len(all_returns), 1)
    min_r, max_r = min(all_returns), max(all_returns)
    padding = 2
    min_r = math.floor(min_r) - padding
    max_r = math.ceil(max_r) + padding
    bin_width = max(1.0, round((max_r - min_r) / 10, 1))
    
    bins = []
    current = min_r
    while current < max_r:
        end = current + bin_width
        count = sum(1 for r in all_returns if current <= r < end)
        bins.append({"start": round(current, 1), "end": round(end, 1), "count": count})
        current = end
    
    return {"bins": bins, "mean": mean}

def compute_qq_plot(monthly_data: dict) -> dict:
    """Compute QQ plot: theoretical vs observed quantiles for monthly returns."""
    matrix = monthly_data.get("matrix", [])
    all_returns = []
    for row in matrix:
        for v in row:
            if v is not None:
                all_returns.append(v)
    
    if len(all_returns) < 10:
        return {"points": []}
    
    import random
    random.seed(42)
    sorted_r = sorted(all_returns)
    n = len(sorted_r)
    mean_r = sum(sorted_r) / n
    # Sample std (using n-1)
    variance = sum((r - mean_r) ** 2 for r in sorted_r) / (n - 1) if n > 1 else 1
    std_r = math.sqrt(max(variance, 1e-10))
    
    points = []
    for i in range(1, n + 1):
        p = i / (n + 1)
        # Approximate inverse normal (Abramowitz & Stegun approximation)
        t = math.sqrt(-2 * math.log(min(p, 1 - p)))
        c0 = 2.515517
        c1 = 0.802853
        c2 = 0.010328
        d1 = 1.432788
        d2 = 0.189269
        d3 = 0.001308
        sign = 1 if p >= 0.5 else -1
        theoretical = sign * (t - (c0 + c1 * t + c2 * t * t) / (1 + d1 * t + d2 * t * t + d3 * t * t * t))
        observed = (sorted_r[i - 1] - mean_r) / std_r
        points.append({
            "theoretical": round(theoretical, 3),
            "observed": round(observed, 3)
        })
    
    return {"points": points}

def compute_rolling_stats(equity_curve: list[dict], window_days: int = 126) -> dict:
    """Compute rolling 6-month (126 trading day) return and volatility."""
    roll = []
    values = [p["value"] for p in equity_curve]
    
    for i in range(window_days, len(values)):
        past = values[i - window_days:i]
        cur_val = values[i]
        prev_val = values[i - window_days]
        
        if prev_val > 0:
            # Rolling return: total return over window, annualized
            roll_ret = ((cur_val / prev_val) - 1)
            # Daily returns for volatility
            daily_rets = [(past[j] / past[j-1]) - 1 for j in range(1, len(past)) if past[j-1] > 0]
            if daily_rets:
                vol = math.sqrt(sum(r * r for r in daily_rets) / len(daily_rets)) * math.sqrt(365)
            else:
                vol = 0
            roll.append({
                "date": equity_curve[i]["date"],
                "rollingReturn": round(roll_ret * 100, 2),
                "rollingVolatility": round(vol * 100, 2)
            })
    
    return {"windowMonths": 6, "series": roll}

def compute_quant_report(strategy_name: str, strategy_id: str, trades: list[dict],
                         start_equity: float = 100.0) -> dict:
    """Compute the full QF-Lib quant report."""
    equity = compute_daily_equity(trades, start_equity)
    monthly = compute_monthly_returns(equity)
    yearly, mean_yearly = compute_yearly_returns(monthly)
    distribution = compute_return_distribution(monthly)
    qq = compute_qq_plot(monthly)
    rolling = compute_rolling_stats(equity)
    
    return {
        "meta": {
            "strategyName": strategy_name,
            "strategyId": strategy_id,
            "generatedAt": datetime.utcnow().isoformat() + "Z",
            "library": "QF-Lib",
            "version": "1.0.0"
        },
        "equityCurve": equity,
        "monthlyReturns": monthly,
        "yearlyReturns": yearly,
        "meanYearlyReturn": mean_yearly,
        "monthlyReturnDistribution": distribution,
        "qqPlot": qq,
        "rollingStats": rolling
    }
