import { useEffect, useMemo, useState } from "react";
import {
  LineChart, Line, XAxis, YAxis, Tooltip, ResponsiveContainer, Legend, CartesianGrid,
  AreaChart, Area, BarChart, Bar
} from "recharts";

// -----------------------------
// Fake JSON example (5 assets, 67 trading days, randomized tickers)
// -----------------------------
function generateFakeJSON() {
  const randTicker = () => Math.random().toString(36).substring(2, 6).toUpperCase();
  const tickers = Array.from({ length: 5 }, () => randTicker());

  const start = new Date("2024-01-02T00:00:00Z");
  const dates = Array.from({ length: 67 }, (_, i) => {
    const d = new Date(start);
    d.setDate(start.getDate() + i);
    return d.toISOString().slice(0, 10);
  });

  const fake: Record<string, { date: string; close: number }[]> = {};
  for (let a = 0; a < 5; a++) {
    const ticker = tickers[a];
    let price = 80 + Math.random() * 40;
    const mu = (Math.random() * 0.1 - 0.05) / 252;
    const sigma = 0.15 + Math.random() * 0.35;
    const series: { date: string; close: number }[] = [];
    for (let i = 0; i < dates.length; i++) {
      if (i > 0) {
        const z = (Math.random() - 0.5) * 1.6 + (Math.random() - 0.5) * 1.6;
        const daily = mu + (sigma / Math.sqrt(252)) * z;
        price *= 1 + daily;
      }
      series.push({ date: dates[i], close: Number(price.toFixed(2)) });
    }
    fake[ticker] = series;
  }
  return fake;
}

const MAX_DAYS = 67;
const maxSteps = 60; // number of picks to reach MAX_DAYS from START_DAY
const START_DAY = 7; // first visible day, first pick occurs at this day

export default function App() {
  const [rawData, setRawData] = useState<Record<string, { date: string; close: number }[]>>({});
  const [assets, setAssets] = useState<string[]>([]);
  const [dates, setDates] = useState<string[]>([]);

  const [step, setStep] = useState(1); // number of picks made
  const [windowLen, setWindowLen] = useState(START_DAY); // initial visible window length
  const [selectedAsset, setSelectedAsset] = useState<string | null>(null);
  const [hoverAsset, setHoverAsset] = useState<string | null>(null);
  const [selections, setSelections] = useState<{ step: number; date: string; asset: string; ret: number }[]>([]);
  const [message, setMessage] = useState("");
  const [confirming, setConfirming] = useState(false);
  const [finalSaved, setFinalSaved] = useState(false);

  // boot with example data
  useEffect(() => {
    if (Object.keys(rawData).length === 0) {
      const example = generateFakeJSON();
      const keys = Object.keys(example);
      const first = example[keys[0]];
      setRawData(example);
      setAssets(keys);
      setDates(first.map((d) => d.date));
      setStep(1);
      setWindowLen(START_DAY);
      setSelections([]);
      setSelectedAsset(null);
      setMessage(`Loaded example: ${keys.length} assets, ${first.length} days.`);
    }
  }, []);

  const isFinal = windowLen >= MAX_DAYS || selections.length >= maxSteps;

  const windowData = useMemo(() => {
    if (!dates.length || windowLen < 2) return [] as any[];
    const sliceDates = dates.slice(0, windowLen);
    return sliceDates.map((date, idx) => {
      const row: Record<string, any> = { date };
      assets.forEach((a) => {
        const base = rawData[a]?.[0]?.close ?? 1;
        const val = rawData[a]?.[idx]?.close ?? base;
        row[a] = base ? val / base : 1;
      });
      return row;
    });
  }, [assets, dates, windowLen, rawData]);

  function realizedNextDayReturn(asset: string) {
    const t = windowLen - 1;
    if (t + 1 >= dates.length) return null as any;
    const series = rawData[asset];
    const ret = series[t + 1].close / series[t].close - 1;
    return { date: dates[t + 1], ret };
  }

  function loadExample() {
    const example = generateFakeJSON();
    const keys = Object.keys(example);
    const first = example[keys[0]];
    setRawData(example);
    setAssets(keys);
    setDates(first.map((d) => d.date));
    setStep(1);
    setWindowLen(START_DAY);
    setSelections([]);
    setSelectedAsset(null);
    setMessage(`Loaded example: ${keys.length} assets, ${first.length} days.`);
    try { localStorage.removeItem("asset_experiment_selections"); } catch {}
  }

  function resetSession() {
    setSelections([]);
    setSelectedAsset(null);
    setStep(1);
    setWindowLen(START_DAY);
    setMessage("Session reset.");
    try { localStorage.removeItem("asset_experiment_selections"); } catch {}
  }

  function onFile(e: any) {
    const f = e.target.files?.[0];
    if (!f) return;
    const reader = new FileReader();
    reader.onload = () => {
      try {
        const json = JSON.parse(String(reader.result));
        const keys = Object.keys(json);
        if (keys.length === 0) throw new Error("Empty dataset");
        const first = json[keys[0]];
        if (!Array.isArray(first) || !first[0]?.date || typeof first[0]?.close !== "number") {
          throw new Error("Invalid series format (need [{date, close}])");
        }
        const ref = new Set(first.map((d: any) => d.date));
        for (const k of keys.slice(1)) {
          for (const p of json[k]) { if (!ref.has(p.date)) throw new Error("Date misalignment across assets"); }
        }
        setRawData(json);
        setAssets(keys);
        setDates(first.map((d: any) => d.date));
        setStep(1);
        setWindowLen(START_DAY);
        setSelections([]);
        setSelectedAsset(null);
        setMessage(`Loaded file: ${keys.length} assets, ${first.length} days.`);
        try { localStorage.removeItem("asset_experiment_selections"); } catch {}
      } catch (err: any) {
        setMessage("Failed to parse JSON: " + err.message);
      }
    };
    reader.readAsText(f);
  }

  function exportLog() {
    const blob = new Blob([JSON.stringify(selections, null, 2)], { type: "application/json" });
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = `selections_${new Date().toISOString().slice(0,10)}.json`;
    a.click();
    URL.revokeObjectURL(url);
  }

  function confirmSelection() {
    if (confirming) return;
    if (!selectedAsset) return setMessage("Select a line first.");
    setConfirming(true);

    const res = realizedNextDayReturn(selectedAsset);
    if (!res) {
      setMessage("No more data available.");
      setConfirming(false);
      return;
    }
    const entry = { step, date: res.date, asset: selectedAsset, ret: res.ret };
    setSelections((prev) => [...prev, entry]);

    setWindowLen((w) => Math.min(w + 1, MAX_DAYS));
    setStep((s) => s + 1);
    setSelectedAsset(null);
    setConfirming(false);
    setMessage(`Pick ${step}: ${selectedAsset} → next-day return ${(res.ret * 100).toFixed(2)}%`);
  }

  const portfolioSeries = useMemo(() => {
    let value = 1;
    const pts = selections.map((s) => {
      value *= 1 + s.ret;
      return { step: s.step, date: s.date, value };
    });
    return [{ step: 0, date: "start", value: 1 }, ...pts];
  }, [selections]);

  const stats = useMemo(() => {
    const rets = selections.map((s) => s.ret);
    const N = rets.length;
    const cum = portfolioSeries.at(-1)?.value ?? 1;
    const mean = N ? rets.reduce((a, b) => a + b, 0) / N : 0;
    const variance = N ? rets.reduce((a, b) => a + (b - mean) ** 2, 0) / N : 0;
    const stdev = Math.sqrt(variance);
    const sharpe = stdev ? (mean * 252) / (stdev * Math.sqrt(252)) : 0;
    const wins = rets.filter((r) => r > 0).length;
    return { cumRet: cum - 1, stdev, sharpe, wins, N };
  }, [portfolioSeries, selections]);
  // ---- Build and auto-save final JSON when finished ----
  function buildFinalPayload() {
    const lastStep = selections.reduce((m, s) => Math.max(m, s.step), 0);
    const start30 = Math.max(1, lastStep - 30 + 1);
    const countsAll = assets.reduce((acc: Record<string, number>, a: string) => { acc[a] = 0; return acc; }, {} as Record<string, number>);
    selections.forEach((s) => { countsAll[s.asset] = (countsAll[s.asset] || 0) + 1; });
    const rankAll = assets.map((a) => ({ asset: a, votes: countsAll[a] || 0 })).sort((x, y) => y.votes - x.votes);

    let value = 1;
    const portfolio = selections.map((s) => { value *= 1 + s.ret; return { step: s.step, date: s.date, value }; });

    const lastCols = Array.from({ length: Math.min(30, lastStep ? lastStep - start30 + 1 : 0) }, (_, i) => start30 + i);
    const heatGrid = assets.map((a) => ({ asset: a, cells: lastCols.map((c) => (selections.some((s) => s.asset === a && s.step === c) ? 1 : 0)) }));

    const rets = selections.map((s) => s.ret);
    const N = rets.length;
    const cum = portfolio.at(-1)?.value ?? 1;
    const mean = N ? rets.reduce((a, b) => a + b, 0) / N : 0;
    const variance = N ? rets.reduce((a, b) => a + (b - mean) ** 2, 0) / N : 0;
    const stdev = Math.sqrt(variance);
    const sharpe = stdev ? (mean * 252) / (stdev * Math.sqrt(252)) : 0;
    const wins = rets.filter((r) => r > 0).length;

    return {
      meta: { saved_at: new Date().toISOString(), start_day: START_DAY, max_days: MAX_DAYS, max_steps: maxSteps },
      assets,
      dates,
      selections,
      portfolio,
      stats: { cumRet: (cum - 1), stdev, sharpe, wins, N },
      preference_all: rankAll,
      heatmap_last30: { cols: lastCols, grid: heatGrid },
    };
  }

  useEffect(() => {
    if (isFinal && !finalSaved) {
      try {
        const payload = buildFinalPayload();
        localStorage.setItem("asset_experiment_final", JSON.stringify(payload));
        const blob = new Blob([JSON.stringify(payload, null, 2)], { type: "application/json" });
        const url = URL.createObjectURL(blob);
        const a = document.createElement("a");
        a.href = url;
        a.download = `run_summary_${new Date().toISOString().slice(0, 10)}.json`;
        a.click();
        URL.revokeObjectURL(url);
        setFinalSaved(true);
      } catch (e) {
        console.warn("Failed to save final JSON:", e);
      }
    }
  }, [isFinal, finalSaved, assets, selections, dates]);

  // Hotkeys: 1..N selects asset, Enter confirms
  useEffect(() => {
    function onKey(e: KeyboardEvent) {
      const tag = (e.target && (e.target as HTMLElement).tagName) || "";
      if (tag === "INPUT" || tag === "TEXTAREA") return;
      const idx = parseInt((e as any).key, 10) - 1;
      if (!Number.isNaN(idx) && idx >= 0 && idx < assets.length) {
        setSelectedAsset(assets[idx]);
      }
      if ((e as any).key === "Enter" && Boolean(selectedAsset) && windowLen < MAX_DAYS) {
        confirmSelection();
      }
    }
    window.addEventListener("keydown", onKey);
    return () => window.removeEventListener("keydown", onKey);
  }, [assets, selectedAsset, windowLen]);

  if (isFinal) {
    return (
      <FinalSummary assets={assets} selections={selections} />
    );
  }

  return (
    <div className="p-6 bg-gray-50 min-h-screen flex flex-col gap-6">
      <div className="flex flex-wrap justify-between items-center gap-3">
        <div className="flex items-center gap-2">
          <h1 className="text-xl font-semibold">Asset Choice Simulation</h1>
          <span className="text-xs text-gray-500">Day {windowLen} / {MAX_DAYS}</span>
        </div>
        <div className="flex items-center gap-2">
          <input type="file" accept="application/json" onChange={onFile} className="text-sm" />
          <button onClick={loadExample} className="text-sm px-3 py-1.5 rounded-xl bg-blue-100 text-blue-800 hover:bg-blue-200">Load Example</button>
          <button onClick={resetSession} className="text-sm px-3 py-1.5 rounded-xl bg-gray-200 hover:bg-gray-300">Reset</button>
          <button onClick={exportLog} className="text-sm px-3 py-1.5 rounded-xl bg-gray-900 text-white hover:bg-black">Export Log</button>
        </div>
      </div>

      <div className="bg-white p-4 rounded-2xl shadow">
        {/* Quick picker (keeps legend & line clicks) */}
        <div className="flex flex-wrap items-center gap-2 mb-3">
          {assets.map((a, i) => (
            <button
              key={`pick-${a}`}
              onClick={() => setSelectedAsset(a)}
              className={`text-xs px-2 py-1 rounded-xl border transition ${selectedAsset===a?"bg-blue-600 text-white border-blue-600":"bg-white text-gray-700 border-gray-200 hover:bg-gray-50"}`}
              title={`Hotkey ${i+1}`}
            >
              <span className="inline-block w-2.5 h-2.5 rounded-full mr-2" style={{backgroundColor: `hsl(${(360/assets.length)*i},70%,50%)`}} />
              {i+1}. {a}
            </button>
          ))}
          {assets.length>0 && (
            <span className="text-xs text-gray-500 ml-1">Hotkeys: 1–{assets.length}, Enter to confirm</span>
          )}
        </div>
        <div className="h-80">
          <ResponsiveContainer width="100%" height="100%">
            <LineChart data={windowData}>
              <CartesianGrid strokeDasharray="3 3" />
              <XAxis dataKey="date" tick={{ fontSize: 10 }} />
              <YAxis domain={["auto", "auto"]} tick={{ fontSize: 10 }} />
              <Tooltip contentStyle={{ fontSize: 12 }} />
              <Legend onClick={(o: any) => setSelectedAsset(o.value)} wrapperStyle={{ cursor: "pointer" }} />
              {assets.map((a, i) => (
                <Line
                  key={a}
                  type="monotone"
                  dataKey={a}
                  strokeWidth={selectedAsset === a ? 5 : hoverAsset === a ? 4 : 2.5}
                  strokeOpacity={selectedAsset && selectedAsset !== a ? 0.3 : 1}
                  dot={false}
                  isAnimationActive={false}
                  stroke={`hsl(${(360 / assets.length) * i},70%,50%)`}
                  onMouseEnter={() => setHoverAsset(a)}
                  onMouseLeave={() => setHoverAsset(null)}
                  onClick={() => setSelectedAsset((p) => (p === a ? null : a))}
                />
              ))}
            </LineChart>
          </ResponsiveContainer>
        </div>

        <div className="flex justify-between items-center mt-3">
          <div className="text-sm text-gray-600">Selected: {selectedAsset ?? "(none)"} {message && <span className="ml-2 text-gray-500">{message}</span>}</div>
          <button
            onClick={confirmSelection}
            disabled={!selectedAsset || windowLen >= MAX_DAYS}
            className={`px-4 py-2 rounded-xl ${selectedAsset && windowLen < MAX_DAYS ? "bg-blue-600 text-white" : "bg-gray-200 text-gray-500"}`}
          >
            Confirm & Next Day →
          </button>
        </div>
      </div>

      <div className="bg-white p-4 rounded-2xl shadow">
        <h2 className="font-medium mb-2">Portfolio</h2>
        <div className="h-64">
          <ResponsiveContainer width="100%" height="100%">
            <AreaChart data={portfolioSeries}>
              <CartesianGrid strokeDasharray="3 3" />
              <XAxis dataKey="step" tick={{ fontSize: 10 }} />
              <YAxis domain={["auto", "auto"]} tick={{ fontSize: 10 }} />
              <Tooltip />
              <Area type="monotone" dataKey="value" stroke="#2563eb" fill="#bfdbfe" />
            </AreaChart>
          </ResponsiveContainer>
        </div>
        <ul className="text-sm text-gray-700 mt-2">
          <li>Cumulative Return: {(stats.cumRet * 100).toFixed(2)}%</li>
          <li>Volatility: {(stats.stdev * 100).toFixed(2)}%</li>
          <li>Sharpe: {stats.sharpe.toFixed(2)}</li>
          <li>Winning Days: {stats.wins}/{stats.N}</li>
        </ul>
      </div>

      {/* Daily selections table */}
      <div className="bg-white p-4 rounded-2xl shadow">
        <h2 className="font-medium mb-2">Daily Selections</h2>
        <div className="overflow-auto rounded-xl border border-gray-100">
          <table className="min-w-full text-xs">
            <thead className="bg-gray-50 text-gray-500">
              <tr>
                <th className="px-2 py-1 text-left">Step</th>
                <th className="px-2 py-1 text-left">Date (t+1)</th>
                <th className="px-2 py-1 text-left">Asset</th>
                <th className="px-2 py-1 text-right">Return</th>
              </tr>
            </thead>
            <tbody>
              {selections.slice().reverse().map((s) => (
                <tr key={`${s.step}-${s.asset}-${s.date}`} className="odd:bg-white even:bg-gray-50">
                  <td className="px-2 py-1">{s.step}</td>
                  <td className="px-2 py-1">{s.date}</td>
                  <td className="px-2 py-1">{s.asset}</td>
                  <td className={`px-2 py-1 text-right ${s.ret>=0?"text-green-600":"text-red-600"}`}>{(s.ret*100).toFixed(2)}%</td>
                </tr>
              ))}
            </tbody>
          </table>
        </div>
      </div>
    </div>
  );
}

// ---------- Final Summary Component ----------
function FinalSummary({ assets, selections }: { assets: string[]; selections: { step: number; date: string; asset: string; ret: number }[] }) {
  // Overall metrics for full run
  const rets = selections.map(s=>s.ret);
  const N = rets.length;
  const cum = rets.reduce((v,r)=> v*(1+r), 1) - 1;
  const mean = N ? rets.reduce((a,b)=>a+b,0)/N : 0;
  const variance = N ? rets.reduce((a,b)=> a + (b-mean)**2, 0)/N : 0;
  const stdev = Math.sqrt(variance);
  const sharpe = stdev ? (mean*252)/(stdev*Math.sqrt(252)) : 0;
  const wins = rets.filter(r=>r>0).length;

  // Preference ranking (ALL picks)
  const countsAll: Record<string, number> = assets.reduce((acc: any,a: string)=>{acc[a]=0;return acc;},{} as Record<string, number>);
  selections.forEach(s=>{ countsAll[s.asset] = (countsAll[s.asset]||0)+1; });
  const rankAll = assets
    .map(a=>({ asset:a, votes: countsAll[a]||0 }))
    .sort((x,y)=> y.votes - x.votes);

  // Heatmap for last 30 steps
  const lastStep = selections.reduce((m,s)=>Math.max(m,s.step),0);
  const start30 = Math.max(1, lastStep - 30 + 1);
  const cols = Array.from({length: Math.min(30,lastStep ? lastStep - start30 + 1 : 0)}, (_,i)=> start30 + i);
  const grid = assets.map(a=>({
    asset:a,
    cells: cols.map(c => selections.some(s=> s.asset===a && s.step===c) ? 1 : 0)
  }));

  return (
    <div className="p-6 bg-gray-50 min-h-screen flex flex-col gap-6">
      <h1 className="text-xl font-semibold">Final Summary</h1>

      <div className="bg-white p-4 rounded-2xl shadow">
        <h2 className="font-medium mb-2">Overall Metrics</h2>
        <ul className="text-sm text-gray-700 space-y-1">
          <li>Total Picks: {N}</li>
          <li>Win Rate: {(N? (wins/N*100):0).toFixed(1)}%</li>
          <li>Cumulative Return: {(cum*100).toFixed(2)}%</li>
          <li>Volatility: {(stdev*100).toFixed(2)}%</li>
          <li>Sharpe (rough): {sharpe.toFixed(2)}</li>
          <li>Top Preference (All): {rankAll[0]?.asset ?? "-"} ({rankAll[0]?.votes ?? 0} votes)</li>
        </ul>
      </div>

      <div className="bg-white p-4 rounded-2xl shadow">
        <h2 className="font-medium mb-2">Selection Preference Ranking (All Assets)</h2>
        <div className="h-56">
          <ResponsiveContainer width="100%" height="100%">
            <BarChart data={rankAll} margin={{ left: 8, right: 8, top: 8, bottom: 8 }}>
              <CartesianGrid strokeDasharray="3 3" />
              <XAxis dataKey="asset" tick={{ fontSize: 10 }} />
              <YAxis allowDecimals={false} tick={{ fontSize: 10 }} />
              <Tooltip />
              <Bar dataKey="votes" fill="#60a5fa" />
            </BarChart>
          </ResponsiveContainer>
        </div>
      </div>

      <div className="bg-white p-4 rounded-2xl shadow">
        <h2 className="font-medium mb-2">Selection Heatmap (Assets × Last 30 Steps)</h2>
        <div className="overflow-auto">
          <table className="text-xs border-collapse">
            <thead>
              <tr>
                <th className="p-1 pr-2 text-left sticky left-0 bg-white">Asset</th>
                {cols.map(c=> (
                  <th key={c} className="px-1 py-1 text-center">{c}</th>
                ))}
              </tr>
            </thead>
            <tbody>
              {grid.map(row => (
                <tr key={row.asset}>
                  <td className="p-1 pr-2 font-medium sticky left-0 bg-white">{row.asset}</td>
                  {row.cells.map((v,j)=> (
                    <td key={j} className="w-6 h-6" style={{
                      background: v? "#2563eb" : "#e5e7eb",
                      opacity: v? 0.9 : 1,
                      border: "1px solid #ffffff"
                    }} />
                  ))}
                </tr>
              ))}
            </tbody>
          </table>
        </div>
      </div>
    </div>
  );
}