// Derived peer groups for Philly schools (K-nearest neighbors on demographics).
//
// Mirrors scripts/loaders/nycenet-comparison-groups.ts (NYC) but adapted:
//   - HARD filter: school_level (Elementary / Middle / High / K-8 / etc.)
//   - HARD filter: admission bucket (collapsed from PhillyAdmissionCategory)
//   - HARD filter: school_type sub-buckets (district/charter — different
//     governance is its own peer universe; the analysis pipeline only
//     compares like-with-like)
//   - WEIGHTED distance: 7 demographic features × Euclidean
//   - K = 40 default; we ALSO derive k=20 and k=30 + emit a flip-diff
//     showing which schools' peer set changes substantively. The 3
//     derivations write to philly_comparison_groups with separate `k` rows;
//     only k=40 is linked into philly_schools.comparison_group_id by default.
//
// Run:  npx tsx scripts/loaders/philly/derive_peer_groups.ts

import crypto from "node:crypto";
import { prisma } from "./_lib";
import type { PhillyAdmissionCategory, PhillySchoolType } from "@prisma/client";

const KS_TO_DERIVE = [20, 30, 40] as const;
const DEFAULT_K = 40;
const YEAR = "2024-25";

const DEMOG_FEATURES = [
  "pct_black",
  "pct_hispanic",
  "pct_asian",
  "pct_white",
  "pct_english_learner",
  "pct_special_ed",
  "pct_econ_disadv",
] as const;
type Feature = (typeof DEMOG_FEATURES)[number];

// Collapse admission categories into broader peer buckets so small universes
// (e.g. 8 special-admission HS) have enough peers.
const ADMISSION_BUCKETS: Record<string, PhillyAdmissionCategory[]> = {
  catchment: ["CATCHMENT"],
  selective: ["SPECIAL_ADMISSION", "CRITERIA_BASED"],
  lottery: ["CITYWIDE_LOTTERY", "CHARTER_LOTTERY"],
  cte: ["CTE"],
  alternative: ["ALTERNATIVE"],
  other: ["OTHER"],
};

function admissionBucket(cat: PhillyAdmissionCategory | null): string {
  if (!cat) return "other";
  for (const [b, cats] of Object.entries(ADMISSION_BUCKETS)) {
    if (cats.includes(cat)) return b;
  }
  return "other";
}

// School-type sub-buckets — district vs charter is its own analytical universe.
// (Per O.2 they're "one universe, tagged" for storage; for peer matching
// they get different buckets so a district school doesn't peer with a charter.)
function typeBucket(t: PhillySchoolType): string {
  switch (t) {
    case "REGULAR":
    case "CONTRACTED":
      return "district";
    case "CHARTER":
      return "charter";
    case "CYBER":
      return "cyber";
    case "COMPCTC":
      return "ctc";
    case "SPECIALED":
      return "specialed";
    default:
      return "other";
  }
}

type SchoolForMatch = {
  ulcsCode: string;
  schoolLevel: string;             // Elementary / Middle / High / K-8 / etc.
  admissionBucket: string;
  typeBucket: string;
  includeInDefaultComparisons: boolean;
  demographics: Partial<Record<Feature, number>>;
};

function demographicDistance(a: SchoolForMatch, b: SchoolForMatch): number {
  let sum = 0;
  let n = 0;
  for (const f of DEMOG_FEATURES) {
    const va = a.demographics[f];
    const vb = b.demographics[f];
    if (va === undefined || vb === undefined) continue;
    const d = va - vb;
    sum += d * d;
    n++;
  }
  if (n === 0) return Infinity;
  return Math.sqrt(sum / n);
}

async function main() {
  console.log(`[derive] loading Philly schools + demographics`);
  const schools = await prisma.phillySchool.findMany({
    select: {
      ulcsCode: true,
      schoolLevel: true,
      schoolType: true,
      admissionType: true,
      includeInDefaultComparisons: true,
      closedAt: true,
    },
  });
  console.log(`[derive]   ${schools.length} total schools`);

  // Demographics are stored as long-format facts in philly_school_year_metrics
  // (subgroup=ALL, metric_key in DEMOG_FEATURES). Pull the latest year per
  // (school, metric).
  const demRows = await prisma.phillySchoolYearMetric.findMany({
    where: {
      metricKey: { in: [...DEMOG_FEATURES] },
      subgroup: "ALL",
      value: { not: null },
    },
    select: { schoolUlcs: true, year: true, metricKey: true, value: true },
    orderBy: [{ schoolUlcs: "asc" }, { year: "desc" }],
  });
  // Map school → metric → latest value
  const demByUlcs = new Map<string, Partial<Record<Feature, number>>>();
  for (const r of demRows) {
    const m = demByUlcs.get(r.schoolUlcs) ?? {};
    const f = r.metricKey as Feature;
    // First-seen wins; orderBy(year desc) → that's the latest
    if (m[f] === undefined) m[f] = r.value!;
    demByUlcs.set(r.schoolUlcs, m);
  }

  const working: SchoolForMatch[] = schools
    .filter((s) => s.closedAt === null && demByUlcs.has(s.ulcsCode))
    .map((s) => ({
      ulcsCode: s.ulcsCode,
      schoolLevel: s.schoolLevel ?? "Unknown",
      admissionBucket: admissionBucket(s.admissionType),
      typeBucket: typeBucket(s.schoolType),
      includeInDefaultComparisons: s.includeInDefaultComparisons,
      demographics: demByUlcs.get(s.ulcsCode)!,
    }));
  console.log(`[derive]   ${working.length} schools with demographics + open`);

  // Group by (school_level, admission_bucket, type_bucket) — the matching universe.
  const byKey = new Map<string, SchoolForMatch[]>();
  for (const s of working) {
    const key = `${s.schoolLevel}|${s.admissionBucket}|${s.typeBucket}`;
    const arr = byKey.get(key) ?? [];
    arr.push(s);
    byKey.set(key, arr);
  }

  console.log(`\n[derive] Matching universes (school_level|admission|type):`);
  const keys = [...byKey.keys()].sort();
  for (const k of keys) {
    console.log(`  ${k.padEnd(50)} ${byKey.get(k)!.length} schools`);
  }

  // Compute ranked peer list per school once (sorted by distance).
  // Then take top-K for each k value.
  type RankedPeer = { ulcsCode: string; d: number };
  const rankedPerSchool = new Map<string, RankedPeer[]>();
  for (const [_key, peers] of byKey) {
    for (const me of peers) {
      const ranked = peers
        .filter((p) => p.ulcsCode !== me.ulcsCode && p.includeInDefaultComparisons)
        .map((p) => ({ ulcsCode: p.ulcsCode, d: demographicDistance(me, p) }))
        .filter((p) => Number.isFinite(p.d))
        .sort((a, b) => a.d - b.d);
      rankedPerSchool.set(me.ulcsCode, ranked);
    }
  }

  // Wipe + recreate comparison groups (idempotent).
  console.log(`\n[derive] resetting philly_comparison_groups...`);
  await prisma.phillySchool.updateMany({ data: { comparisonGroupId: null } });
  await prisma.phillyComparisonGroup.deleteMany({});

  // For each k, derive groups and write rows. Keep mapping (k → ulcs → groupId)
  // so we can flip-diff and link the default k.
  type GroupRow = {
    id: string;
    k: number;
    ulcs: string;
    methodologyNotes: { k: number; year: string; school_level: string; admission_bucket: string; type_bucket: string; peers: string[] };
  };
  const groupsByK: Record<number, Map<string, GroupRow>> = {};
  for (const k of KS_TO_DERIVE) {
    groupsByK[k] = new Map();
    let dropped = 0;
    for (const me of working) {
      const ranked = rankedPerSchool.get(me.ulcsCode) ?? [];
      const top = ranked.slice(0, k);
      if (top.length < 5) { dropped++; continue; }
      const id = crypto.randomUUID();
      groupsByK[k].set(me.ulcsCode, {
        id,
        k,
        ulcs: me.ulcsCode,
        methodologyNotes: {
          k,
          year: YEAR,
          school_level: me.schoolLevel,
          admission_bucket: me.admissionBucket,
          type_bucket: me.typeBucket,
          peers: top.map((t) => t.ulcsCode),
        },
      });
    }
    console.log(`[derive]   k=${k}: ${groupsByK[k].size} groups (${dropped} schools dropped with <5 peers)`);
  }

  // Bulk-insert all groups for all k.
  console.log(`\n[derive] inserting comparison_groups rows...`);
  const allGroups = Object.values(groupsByK).flatMap((m) => [...m.values()]);
  for (let i = 0; i < allGroups.length; i += 500) {
    const batch = allGroups.slice(i, i + 500);
    await prisma.phillyComparisonGroup.createMany({
      data: batch.map((g) => ({
        id: g.id,
        year: YEAR,
        k: g.k,
        description: `KNN-${g.k} on demographics; hard filters school_level + admission + type bucket`,
        methodologyNotes: g.methodologyNotes as unknown as object,
      })),
    });
  }

  // Link only the DEFAULT_K group on each PhillySchool.
  console.log(`[derive] linking default k=${DEFAULT_K} into philly_schools.comparison_group_id...`);
  const defaultGroups = groupsByK[DEFAULT_K];
  for (const [ulcs, g] of defaultGroups) {
    await prisma.phillySchool.update({
      where: { ulcsCode: ulcs },
      data: { comparisonGroupId: g.id },
    });
  }
  console.log(`[derive]   linked ${defaultGroups.size} schools`);

  // FLIP DIFF: schools whose top-20 peers under k=20 vs k=40 differ substantively.
  console.log(`\n[derive] FLIP DIFF (k=20 vs k=40):`);
  let nDifferentTop10 = 0;
  for (const [ulcs, g20] of groupsByK[20]) {
    const g40 = groupsByK[40].get(ulcs);
    if (!g40) continue;
    const top10_20 = new Set(g20.methodologyNotes.peers.slice(0, 10));
    const top10_40 = new Set(g40.methodologyNotes.peers.slice(0, 10));
    let common = 0;
    for (const p of top10_20) if (top10_40.has(p)) common++;
    if (common < 8) nDifferentTop10++;
  }
  console.log(`[derive]   ${nDifferentTop10} schools where the top-10 peer set differs by ≥3 members between k=20 and k=40`);
  console.log(`[derive]   (≤2-member difference is typical; this counts the substantive flips)`);
  console.log(`\n[derive] done — ${allGroups.length} groups across k=${KS_TO_DERIVE.join(",")}`);
}

main()
  .catch((e) => { console.error(e); process.exit(1); })
  .finally(() => prisma.$disconnect());