import { createMemo, createSignal, For, Show, JSXElement } from "solid-js";
import { US_STATES, US_VIEWBOX } from "./usStates.ts";

// A choropleth of the 50 states + DC, plus optional lat/lng markers. Like uikit/Chart it
// draws SVG as an innerHTML string (the Go Solid compiler will not namespace control-flow
// SVG — see Chart.tsx). The state boundaries are pre-projected with d3's albersUsa into a
// 960×600 box (usStates.ts); the SAME projection is reimplemented below so that lat/lng
// points land exactly on top of the states. It is a faithful, dependency-free port —
// validated to 0px against d3-geo, Alaska and Hawaii insets included.
//
// The map scales by viewBox rather than by measurement (there is no axis text to keep
// crisp), so no ResizeObserver: viewBox 960×600 + a 960/600 aspect-ratio box fills the
// column. Hover reads e.target's data-attributes — the specific state path or point marker
// under the pointer — so no coordinate maths is needed to know what is being pointed at.

// ── the albersUsa projection (ported from d3-geo, scale 1280, translate [480,300]) ──────

const RAD = Math.PI / 180, TAU = 2 * Math.PI;
function conicEqualAreaRaw(y0: number, y1: number) {
    const sy0 = Math.sin(y0), n = (sy0 + Math.sin(y1)) / 2;
    const c = 1 + sy0 * (2 * n - sy0), r0 = Math.sqrt(c) / n;
    return (lambda: number, phi: number): [number, number] => {
        const r = Math.sqrt(c - 2 * n * Math.sin(phi)) / n;
        return [r * Math.sin(lambda * n), r0 - r * Math.cos(lambda * n)];
    };
}
// One conic-equal-area lobe. `center` is given in the rotated frame (near 0° lon), so it is
// not re-rotated; the input point is rotated by `rotateLon` before projecting.
function albersLobe(rotateLon: number, centerLon: number, centerLat: number, p0: number, p1: number, scale: number, tx: number, ty: number) {
    const raw = conicEqualAreaRaw(p0 * RAD, p1 * RAD);
    const rot = (lon: number) => { const l = (lon + rotateLon) * RAD; return ((l + Math.PI) % TAU + TAU) % TAU - Math.PI; };
    const [cx, cy] = raw(centerLon * RAD, centerLat * RAD);
    return (lon: number, lat: number): [number, number] => {
        const [x, y] = raw(rot(lon), lat * RAD);
        return [tx + scale * (x - cx), ty - scale * (y - cy)];
    };
}
const K = 1280, TX = 480, TY = 300, EPS = 1e-6;
const _lower48 = albersLobe(96, -0.6, 38.7, 29.5, 45.5, K, TX, TY);
const _alaska = albersLobe(154, -2, 58.5, 55, 65, K * 0.35, TX - 0.307 * K, TY + 0.201 * K);
const _hawaii = albersLobe(157, -3, 19.9, 8, 18, K, TX - 0.205 * K, TY + 0.212 * K);
const inBox = (p: [number, number], x0: number, y0: number, x1: number, y1: number) =>
    p[0] >= x0 && p[0] <= x1 && p[1] >= y0 && p[1] <= y1;

// Project [lng, lat] to the 960×600 map, choosing the lower-48 / Alaska / Hawaii lobe the
// way albersUsa does — by which one's clip box the point falls in. null if off-map.
export function projectUS(lng: number, lat: number): [number, number] | null {
    let p = _lower48(lng, lat);
    if (inBox(p, TX - 0.455 * K, TY - 0.238 * K, TX + 0.455 * K, TY + 0.238 * K)) return p;
    p = _alaska(lng, lat);
    if (inBox(p, TX - 0.425 * K + EPS, TY + 0.120 * K + EPS, TX - 0.214 * K - EPS, TY + 0.234 * K - EPS)) return p;
    p = _hawaii(lng, lat);
    if (inBox(p, TX - 0.214 * K + EPS, TY + 0.166 * K + EPS, TX - 0.115 * K - EPS, TY + 0.234 * K - EPS)) return p;
    return null;
}

// ── the component ────────────────────────────────────────────────────────────────

export interface USHeatmapPoint {
    lat: number;
    lng: number;
    value?: number;
    label?: string;
}

export interface USHeatmapProps {
    // State value map: USPS code (e.g. "CA", "TX", "DC") → number. States present are
    // shaded on the sequential ramp; states absent are drawn in the no-data neutral.
    data?: Record<string, number>;
    // lat/lng markers, projected onto the map. Points outside the US are dropped.
    points?: USHeatmapPoint[];

    height?: number; // px; omit to size from the container width (960:600 aspect).
    class?: string;
    steps?: number;  // choropleth buckets, 1–6 (default 6, the token count).
    tooltip?: boolean;
    valueFormat?: (v: number) => string;
    pointColor?: string;   // default var(--color-chart-1) (blue).
    pointRadius?: number;  // fixed dot radius (default 5); the MAXIMUM radius when proportional.
    pointOpacity?: number; // dot fill opacity (default 0.85), so the state beneath still reads.
    // Scale each dot's AREA by its value (radius ∝ √value) so a bigger dot means "more" —
    // area, not radius, because the eye reads a circle by its area.
    proportional?: boolean;
    // Override a state's tooltip name (default the built-in full name).
    stateName?: (code: string) => string;
}

const CHOROPLETH_STEPS = 6; // must match --color-choropleth-1..N in theme.css
const POINT_OPACITY = 0.85; // dots are slightly see-through so the state beneath still reads
const _intl = () => new Intl.NumberFormat("en-US", { maximumFractionDigits: 2 });
let _fmt: Intl.NumberFormat | null = null;
const defaultFormat = (v: number) => (Number.isFinite(v) ? (_fmt ??= _intl()).format(v) : String(v));
const clamp = (v: number, lo: number, hi: number) => Math.min(hi, Math.max(lo, v));
const esc = (s: unknown) =>
    String(s).replace(/[&<>"]/g, (c) => (c === "&" ? "&amp;" : c === "<" ? "&lt;" : c === ">" ? "&gt;" : "&quot;"));

type Hover =
    | { kind: "state"; code: string }
    | { kind: "point"; idx: number }
    | null;

export function USHeatmap(props: USHeatmapProps): JSXElement {
    let wrap: HTMLDivElement | undefined;
    const [hover, setHover] = createSignal<Hover>(null);
    const [pointer, setPointer] = createSignal<[number, number]>([0, 0]);

    const fmt = (v: number) => (props.valueFormat ?? defaultFormat)(v);
    const steps = () => clamp(props.steps ?? CHOROPLETH_STEPS, 1, CHOROPLETH_STEPS);
    const stateName = (code: string) => (props.stateName ? props.stateName(code) : US_STATES[code]?.name ?? code);

    // The value range across the states that have data, for quantising into ramp buckets.
    const range = createMemo(() => {
        const vals = Object.values(props.data ?? {}).filter((v) => Number.isFinite(v));
        return vals.length ? { min: Math.min(...vals), max: Math.max(...vals), has: true } : { min: 0, max: 0, has: false };
    });
    const bucket = (v: number) => {
        const r = range();
        const t = r.max > r.min ? (v - r.min) / (r.max - r.min) : 1;
        return clamp(Math.floor(t * steps()), 0, steps() - 1) + 1;
    };

    // projected markers (drop anything off-map), kept with their original index for hover.
    const points = createMemo(() =>
        (props.points ?? []).map((pt, idx) => ({ pt, idx, xy: projectUS(pt.lng, pt.lat) }))
            .filter((m): m is { pt: USHeatmapPoint; idx: number; xy: [number, number] } => m.xy !== null));

    const body = createMemo(() => {
        const data = props.data ?? {};
        const hv = hover();
        const proportional = !!props.proportional;
        const maxR = props.pointRadius ?? (proportional ? 16 : 5);
        const minR = Math.min(3, maxR * 0.35);
        const maxV = proportional ? Math.max(1, ...points().map((m) => Math.max(0, m.pt.value ?? 0))) : 1;
        const radiusOf = (v: number | undefined) =>
            proportional ? minR + (maxR - minR) * Math.sqrt(clamp((v ?? 0) / maxV, 0, 1)) : maxR;
        const pc = esc(props.pointColor ?? "var(--color-chart-1)");
        const out: string[] = [];

        for (const code in US_STATES) {
            const st = US_STATES[code];
            const has = Object.prototype.hasOwnProperty.call(data, code) && Number.isFinite(data[code]);
            const fill = has ? `var(--color-choropleth-${bucket(data[code])})` : "var(--color-surface-strong)";
            const isHover = hv?.kind === "state" && hv.code === code;
            out.push(`<path d="${st.d}" data-state="${code}" fill="${fill}" fill-opacity="${isHover ? 0.82 : 1}" stroke="var(--color-surface)" stroke-width="0.8"/>`);
        }

        for (const m of points()) {
            const isHover = hv?.kind === "point" && hv.idx === m.idx;
            const r = radiusOf(m.pt.value);
            out.push(`<circle data-pt="${m.idx}" cx="${m.xy[0]}" cy="${m.xy[1]}" r="${isHover ? r + 2 : r}" fill="${pc}" fill-opacity="${props.pointOpacity ?? POINT_OPACITY}"/>`);
        }
        return out.join("");
    });

    const onMove = (e: PointerEvent) => {
        if (props.tooltip === false) return;
        const t = e.target as Element;
        const pIdx = t.getAttribute?.("data-pt");
        const code = t.getAttribute?.("data-state");
        if (pIdx != null) setHover({ kind: "point", idx: +pIdx });
        else if (code != null) setHover({ kind: "state", code });
        else setHover(null);
        if (wrap) {
            const r = wrap.getBoundingClientRect();
            setPointer([e.clientX - r.left, e.clientY - r.top]);
        }
    };

    const tip = createMemo(() => {
        const hv = hover();
        if (!hv) return null;
        if (hv.kind === "point") {
            const pt = (props.points ?? [])[hv.idx];
            if (!pt) return null;
            return {
                title: pt.label ?? `${pt.lat.toFixed(2)}, ${pt.lng.toFixed(2)}`,
                value: pt.value != null ? fmt(pt.value) : "",
                swatch: props.pointColor ?? "var(--color-chart-1)",
            };
        }
        const v = (props.data ?? {})[hv.code];
        const has = v != null && Number.isFinite(v);
        return {
            title: stateName(hv.code),
            value: has ? fmt(v) : "no data",
            swatch: has ? `var(--color-choropleth-${bucket(v)})` : "var(--color-surface-strong)",
        };
    });

    return (
        <div ref={wrap} class={"relative w-full" + (props.class ? " " + props.class : "")}>
            <svg viewBox={US_VIEWBOX} role="img" class="block w-full"
                style={props.height ? { height: `${props.height}px` } : { "aspect-ratio": "960 / 600" }}
                innerHTML={body()} onpointermove={onMove} onpointerleave={() => setHover(null)} />

            <Show when={props.tooltip !== false && tip()}>
                {(t) => (
                    <div class="pointer-events-none absolute z-10 min-w-28 max-w-64 rounded-default border border-line bg-surface px-3 py-2 text-xs shadow-lg"
                        style={{
                            left: `${clamp(pointer()[0], 8, 100000)}px`,
                            top: `${Math.max(8, pointer()[1])}px`,
                            transform: "translate(-50%, calc(-100% - 12px))",
                        }}>
                        <div class="flex items-center gap-2">
                            <span class="inline-block h-2.5 w-2.5 shrink-0 rounded-xs" style={{ "background-color": t().swatch }} />
                            <span class="font-medium text-ink">{t().title}</span>
                        </div>
                        <Show when={t().value}>
                            <div class="mt-1 font-semibold text-ink" style={{ "font-variant-numeric": "tabular-nums" }}>{t().value}</div>
                        </Show>
                    </div>
                )}
            </Show>

            <Show when={range().has}>
                <ChoroplethLegend min={range().min} max={range().max} steps={steps()} fmt={fmt} />
            </Show>
        </div>
    );
}

function ChoroplethLegend(p: { min: number; max: number; steps: number; fmt: (v: number) => string }): JSXElement {
    const swatches = () => Array.from({ length: p.steps }, (_, i) => i + 1);
    return (
        <div class="mt-3 flex items-center gap-2 text-xs text-ink-muted">
            <span style={{ "font-variant-numeric": "tabular-nums" }}>{p.fmt(p.min)}</span>
            <div class="flex overflow-hidden rounded-xs">
                <For each={swatches()}>{(k) => (
                    <span class="h-3 w-6" style={{ "background-color": `var(--color-choropleth-${k})` }} />
                )}</For>
            </div>
            <span style={{ "font-variant-numeric": "tabular-nums" }}>{p.fmt(p.max)}</span>
        </div>
    );
}
