"""
PoliticoAI – Hugging Face Space edition
Light-weight API that mirrors the full contract from the Production Guide.
Data: small synthetic sample so the demo works instantly.
"""
from fastapi import FastAPI, HTTPException, UploadFile, File
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import pandas as pd
import numpy as np
from typing import List, Dict, Any

# ------------------------------------------------------------------
# 1. tiny synthetic dataset (mirrors the schema in the Guide)
# ------------------------------------------------------------------
PRECINCT_STORE: List[Dict[str, Any]] = [
    {
        "id": "P01",
        "name": "Urban Ward 1",
        "county": "Demo County",
        "classification": "Urban",
        "registered": 3500,
        "latest_d_pct": 63.0,
        "base_turnout": 0.70,
        "results": {
            2020: {"d_pct": 65, "turnout": 72},
            2022: {"d_pct": 60, "turnout": 50},
            2024: {"d_pct": 63, "turnout": 70},
        },
        "demographics": {
            "pct_college": 68,
            "pct_under_35": 40,
            "pct_35_64": 45,
            "pct_65plus": 15,
            "median_income": 48000,
            "pct_white": 55,
            "pct_black": 20,
            "pct_hispanic": 20,
            "pct_asian": 5,
        },
        "centroid": [-75.47, 40.61],
    }
    for i in range(1, 11)  # 10 precincts
]

for p in PRECINCT_STORE:
    p["id"] = f"P{i:02d}"
    p["name"] = f"Precinct {i:02d}"
    p["latest_d_pct"] = 45 + np.random.normal(0, 5)  # competitive
    p["base_turnout"] = 0.65 + np.random.normal(0, 0.05)

# ------------------------------------------------------------------
# 2. FastAPI boiler-plate
# ------------------------------------------------------------------
app = FastAPI(title="PoliticoAI API", version="1.0")
app.mount("/", StaticFiles(directory="static", html=True), name="static")

# ------------------------------------------------------------------
# 3. Models (match the Production Guide contract)
# ------------------------------------------------------------------
class SimulationParams(BaseModel):
    iterations: int = 10000
    environment: str = "neutral"
    polling_error: float = 3.5
    turnout_variability: float = 5.0

class SegmentParams(BaseModel):
    primary: str = "age"
    secondary: str = "none"
    metric: str = "size"

# ------------------------------------------------------------------
# 4. API endpoints (exact contract from Guide appendix)
# ------------------------------------------------------------------
@app.get("/api/health")
def health():
    return {"status": "ok"}

@app.get("/api/precincts")
def get_precincts():
    # strip heavy geometry if present, send only numbers
    return [
        {
            "id": p["id"],
            "name": p["name"],
            "classification": p["classification"],
            "registered_voters": p["registered"],
            "latest_d_pct": p["latest_d_pct"],
            "base_turnout": p["base_turnout"],
            "demographics": p["demographics"],
        }
        for p in PRECINCT_STORE
    ]

@app.post("/api/simulation/montecarlo")
def run_monte_carlo(params: SimulationParams):
    env_map = {"neutral": 0, "d3": 3, "r3": -3, "d5": 5, "r5": -5}
    env_shift = env_map.get(params.environment, 0)
    margins = []
    for _ in range(params.iterations):
        env_noise = np.random.normal(env_shift, params.polling_error)
        total_d, total_r = 0, 0
        for pct in PRECINCT_STORE:
            base_d = pct["latest_d_pct"] / 100
            adj_d = np.clip(base_d + env_noise / 100, 0.05, 0.95)
            turnout = np.clip(
                np.random.normal(pct["base_turnout"], params.turnout_variability / 100),
                0.2, 0.95,
            )
            votes = int(pct["registered"] * turnout)
            total_d += int(votes * adj_d)
            total_r += votes - int(votes * adj_d)
        margins.append((total_d - total_r) / (total_d + total_r) * 100)
    margins = np.array(margins)
    win_prob = float((margins > 0).mean() * 100)
    return {
        "win_prob": win_prob,
        "expected_margin": float(margins.mean()),
        "recount_prob": float((np.abs(margins) < 0.5).mean() * 100),
        "margins": margins.tolist(),
        "ci_95": [float(np.percentile(margins, 2.5)), float(np.percentile(margins, 97.5))],
    }

@app.post("/api/demographics/segment")
def segment_demographics(params: SegmentParams):
    # lightweight demo segmentation
    segments = ["Under 35", "35-64", "65+"]
    treemap = {
        "labels": ["District"] + segments,
        "parents": [""] + ["District"] * len(segments),
        "values": [3500, 4000, 2500],
    }
    matrix = [
        {"segment": f"Age: {seg}", "size": 3500 if seg == "Under 35" else 4000 if seg == "35-64" else 2500,
         "turnout": "69%", "lean": "D+2", "persuadability": "High", "strategy": "Persuade"}
        for seg in segments
    ]
    return {"treemap": treemap, "targeting_matrix": matrix}

@app.post("/api/precincts/analyze")
def analyze_precincts():
    # minimal clustering demo
    metrics = []
    for p in PRECINCT_STORE:
        pvi = p["latest_d_pct"] - 50
        swing = np.random.uniform(1, 8)
        to_delta = np.random.uniform(15, 30)
        net = (swing + to_delta) * p["registered"] / 1000
        metrics.append(
            {
                "id": p["id"],
                "name": p["name"],
                "pvi": round(pvi, 1),
                "swingScore": round(swing, 1),
                "turnoutDelta": round(to_delta, 1),
                "net": round(net),
                "tier": "High Swing" if swing > 4.5 else "Blue Base" if pvi > 0 else "Red Lean",
            }
        )
    metrics = sorted(metrics, key=lambda x: x["net"], reverse=True)
    return {"scatter_data": metrics, "ranking": metrics[:15], "table": metrics[:20]}

@app.post("/api/upload/results")
async def upload_results(file: UploadFile = File(...)):
    # accept any CSV, store in memory (or persist to disk if you like)
    df = pd.read_csv(file.file)
    # here you would validate, clean, merge into PRECINCT_STORE, etc.
    return {"rows": len(df), "columns": list(df.columns)}

# ------------------------------------------------------------------
# 5. CORS — Spaces needs this for browser access
# ------------------------------------------------------------------
from fastapi.middleware.cors import CORSMiddleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # restrict in prod
    allow_methods=["*"],
    allow_headers=["*"],
)

# ------------------------------------------------------------------
# 6. Entry-point guard (lets uvicorn import without running)
# ------------------------------------------------------------------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)