#!/usr/bin/env python3

import argparse
import os
import sys
import json
import math
from datetime import datetime

# Optional DB import
try:
    import test_pensav_db  # your separate DB module
except ImportError:
    test_pensav_db = None

# Constants
DEFAULT_ALLOCS = {"Equities": 0.80, "Bonds": 0.15, "Cash": 0.05}
DEFAULT_GROWTH = {"Equities": 0.06, "Bonds": 0.03, "Cash": 0.005}
SCENARIOS = [
    ("Soft correction", 3, 1.0, 0.35, {"Equities": -0.10, "Bonds": 0.00, "Cash": 0.00}),
    ("Base (no big shock)", 0, 0.0, 0.30, {"Equities": 0.00, "Bonds": 0.00, "Cash": 0.00}),
    ("Bear market", 6, 2.0, 0.25, {"Equities": -0.25, "Bonds": -0.05, "Cash": 0.00}),
    ("Crash", 9, 4.0, 0.10, {"Equities": -0.40, "Bonds": -0.10, "Cash": 0.00}),
]
HORIZON_YEARS = 5
MONTHS = HORIZON_YEARS * 12

def ann_to_monthly(r):
    return (1 + r)**(1/12) - 1

def build_baseline(start_value, allocs, monthly_growth):
    baseline = [start_value]
    for t in range(1, MONTHS + 1):
        step = sum(baseline[-1] * w * monthly_growth[a] for a, w in allocs.items())
        baseline.append(baseline[-1] + step)
    return baseline

def build_paths(start_value, baseline, allocs):
    paths = {}
    end_values = []
    probs = []

    for name, m_trough, y_recover, prob, shock in SCENARIOS:
        path = [float('nan')] * (MONTHS + 1)
        path[0] = start_value

        if name == "Base (no big shock)":
            path = baseline[:]
        else:
            trough_val = sum(start_value * w * (1 + shock[a]) for a, w in allocs.items())

            # Decline to trough
            for t in range(1, m_trough + 1):
                frac = t / m_trough
                path[t] = start_value + frac * (trough_val - start_value)

            # Recover to baseline
            rec_months = max(int(round(y_recover * 12)), 1)
            for t in range(m_trough + 1, m_trough + rec_months + 1):
                if t > MONTHS: break
                frac = (t - m_trough) / rec_months
                target = baseline[t]
                path[t] = trough_val + frac * (target - trough_val)

            # Follow baseline
            last = max(i for i, v in enumerate(path) if not math.isnan(v))
            for t in range(last + 1, MONTHS + 1):
                path[t] = baseline[t]

        paths[name] = path
        end_values.append(path[-1])
        probs.append(prob)

    return paths, end_values, probs

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--out', required=True, help='Output JSON path')
    parser.add_argument('--use-db', action='store_true', help='Enable DB mode')
    parser.add_argument('--start-value', type=float, default=500000, help='Initial portfolio value')
    args = parser.parse_args()

    # Defaults
    start_value = args.start_value
    allocs = DEFAULT_ALLOCS
    growth = DEFAULT_GROWTH

    # Optional DB override
    if args.use_db:
        if not test_pensav_db:
            print("Error: DB mode requested but test_pensav_db module not found.")
            sys.exit(1)
        try:
            db_data = test_pensav_db.fetch_data()
            print("✅ DB data loaded:", db_data)
            start_value = db_data.get("start_value", start_value)
            allocs = db_data.get("allocations", allocs)
            growth = db_data.get("growth", growth)
        except Exception as e:
            print("❌ DB error:", e)
            sys.exit(1)

    monthly_growth = {k: ann_to_monthly(v) for k, v in growth.items()}
    baseline = build_baseline(start_value, allocs, monthly_growth)
    paths, end_values, probs = build_paths(start_value, baseline, allocs)

    ps = [p / sum(probs) for p in probs]
    exp_end = sum(v * w for v, w in zip(end_values, ps))

    out = {
        "as_of": datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC"),
        "start_value": start_value,
        "months": MONTHS,
        "baseline": baseline,
        "paths": paths,
        "expected_end_value": exp_end,
        "allocations": allocs,
        "growth": growth,
        "scenarios": [
            {"name": n, "months_to_trough": m, "years_to_recover": y, "prob": p, "shock": s}
            for (n, m, y, p, s) in SCENARIOS
        ],
    }

    with open(args.out, 'w') as f:
        json.dump(out, f)

    print(f"✅ Wrote {args.out} ({len(out['baseline'])} points)")

if __name__ == "__main__":
    main()