/**
 * llmindex.io — index refit: response matrices → Python 2PL/BT fit → score run
 * Author:  Simon-Pierre Boucher
 * Contact: contact@spboucher.ai
 * License: Proprietary — © Simon-Pierre Boucher, all rights reserved
 */
import { spawnSync } from 'node:child_process';
import { createHash } from 'node:crypto';
import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs';
import { dirname, join } from 'node:path';
import { fileURLToPath } from 'node:url';
import { prisma } from '@llmindex/db';
import {
  DUEL_DOMAINS,
  GLOBAL_DOMAIN,
  INDEX_VERSION,
  IRT_DOMAINS,
  IRT_HYPERPARAMS,
  domainScore,
  globalIndex,
  type Domain,
  type ScoreWithCI,
} from '@llmindex/scoring';

const REPO_ROOT = join(dirname(fileURLToPath(import.meta.url)), '..', '..', '..');
const RUNS_DIR = join(REPO_ROOT, 'data', 'runs');
const FIT_SCRIPT = join(REPO_ROOT, 'apps', 'psychometrics', 'fit.py');
// Prefer the psychometrics venv (numpy) when present; fall back to system python3.
const VENV_PYTHON = join(REPO_ROOT, 'apps', 'psychometrics', '.venv', 'bin', 'python3');
const PYTHON = existsSync(VENV_PYTHON) ? VENV_PYTHON : 'python3';

const MIN_MODELS = 2;
const MIN_ITEMS = 10;

interface DomainMatrix {
  models: string[]; // slugs
  items: string[]; // eval_item ids
  is_anchor: boolean[];
  /** responses[m][i] ∈ 0 | 1 | null */
  responses: (0 | 1 | null)[][];
}

interface FitOutput {
  domains: Record<
    string,
    {
      models: Array<{ slug: string; theta: number; se: number }>;
      items: Array<{ id: string; a: number; b: number }>;
      diagnostics: Record<string, unknown>;
    }
  >;
  duel_domains?: Record<
    string,
    {
      models: Array<{ slug: string; theta: number; se: number }>;
      diagnostics: Record<string, unknown>;
    }
  >;
}

interface DuelMatrix {
  models: string[];
  /** wins[i][j] = wins of i over j (ties pre-encoded as 0.5 each). */
  wins: number[][];
}

interface AuxMetrics {
  consistency: number | null;
  calibration: number | null;
  contaminationDelta: number | null;
  latencyP50: number | null;
  costPer1kItems: number | null;
}

function percentile(values: number[], p: number): number | null {
  if (values.length === 0) return null;
  const sorted = [...values].sort((a, b) => a - b);
  return sorted[Math.min(sorted.length - 1, Math.floor(p * sorted.length))]!;
}

/** Expected calibration error over 10 confidence bins → calibration = 1 − ECE. */
function calibrationFrom(pairs: Array<{ confidence: number; correct: boolean }>): number | null {
  if (pairs.length < 10) return null;
  const bins = Array.from({ length: 10 }, () => ({ n: 0, conf: 0, acc: 0 }));
  for (const p of pairs) {
    const b = bins[Math.min(9, Math.floor(p.confidence * 10))]!;
    b.n += 1;
    b.conf += p.confidence;
    b.acc += p.correct ? 1 : 0;
  }
  let ece = 0;
  for (const b of bins) {
    if (b.n === 0) continue;
    ece += (b.n / pairs.length) * Math.abs(b.acc / b.n - b.conf / b.n);
  }
  return 1 - ece;
}

export async function runRefit(): Promise<{ runId: string } | null> {
  // Degraded batches are excluded from ranking until re-run (§6); only
  // batches of the CURRENT index version enter a fit (methodology coherence).
  const batches = await prisma.scoreRun.findMany({
    where: { kind: 'eval_batch', status: 'complete', indexVersion: INDEX_VERSION },
    select: { id: true },
  });
  const duelBatches = await prisma.scoreRun.findMany({
    where: { kind: 'duel_batch', status: 'complete', indexVersion: INDEX_VERSION },
    select: { id: true },
  });
  if (batches.length === 0 && duelBatches.length === 0) {
    console.warn('[refit] no complete batches for this index version — nothing to fit');
    return null;
  }
  const batchIds = batches.map((b) => b.id);

  const matrices: Record<string, DomainMatrix> = {};
  const aux: Record<string, Record<string, AuxMetrics>> = {}; // domain → slug → metrics

  for (const domain of IRT_DOMAINS) {
    const responses = await prisma.modelResponse.findMany({
      where: {
        runId: { in: batchIds },
        sampleIndex: 0,
        correct: { not: null },
        item: { domain },
      },
      include: {
        model: { select: { slug: true } },
        item: { select: { id: true, isAnchor: true } },
      },
      orderBy: { createdAt: 'asc' },
    });
    if (responses.length === 0) continue;

    // Latest response per model×item wins (re-runs supersede).
    const latest = new Map<string, (typeof responses)[number]>();
    for (const r of responses) latest.set(`${r.model.slug}::${r.item.id}`, r);

    const modelSlugs = [...new Set([...latest.values()].map((r) => r.model.slug))].sort();
    const itemIds = [...new Set([...latest.values()].map((r) => r.item.id))].sort();
    if (modelSlugs.length < MIN_MODELS || itemIds.length < MIN_ITEMS) {
      console.warn(
        `[refit] domain ${domain}: ${modelSlugs.length} models × ${itemIds.length} items — below minimum, skipped`,
      );
      continue;
    }
    const anchorSet = new Map(
      [...latest.values()].map((r) => [r.item.id, r.item.isAnchor] as const),
    );
    matrices[domain] = {
      models: modelSlugs,
      items: itemIds,
      is_anchor: itemIds.map((id) => anchorSet.get(id) ?? false),
      responses: modelSlugs.map((slug) =>
        itemIds.map((itemId) => {
          const r = latest.get(`${slug}::${itemId}`);
          return r ? ((r.correct ? 1 : 0) as 0 | 1) : null;
        }),
      ),
    };

    // Auxiliary sub-metrics per model for this domain.
    aux[domain] = {};
    for (const slug of modelSlugs) {
      const mine = [...latest.values()].filter((r) => r.model.slug === slug);
      const anchorGraded = mine.filter((r) => r.item.isAnchor);
      const freshGraded = mine.filter((r) => !r.item.isAnchor);
      const acc = (rs: typeof mine): number | null =>
        rs.length ? rs.filter((r) => r.correct).length / rs.length : null;
      const anchorAcc = acc(anchorGraded);
      const freshAcc = acc(freshGraded);

      // Consistency: same item, k samples — fraction agreeing with the modal answer.
      const kSamples = await prisma.modelResponse.findMany({
        where: {
          runId: { in: batchIds },
          model: { slug },
          item: { domain },
          answerExtracted: { not: null },
        },
        select: { itemId: true, answerExtracted: true },
      });
      const byItem = new Map<string, string[]>();
      for (const s of kSamples) {
        byItem.set(s.itemId, [...(byItem.get(s.itemId) ?? []), s.answerExtracted!]);
      }
      const consistencies: number[] = [];
      for (const answers of byItem.values()) {
        if (answers.length < 2) continue;
        const counts = new Map<string, number>();
        for (const a of answers) counts.set(a, (counts.get(a) ?? 0) + 1);
        consistencies.push(Math.max(...counts.values()) / answers.length);
      }

      const confPairs = mine
        .filter((r) => r.confidence != null)
        .map((r) => ({ confidence: r.confidence!, correct: r.correct === true }));
      const latencies = mine.map((r) => r.latencyMs).filter((v): v is number => v != null);
      const costs = mine.map((r) => r.costUsd).filter((v): v is number => v != null);

      aux[domain][slug] = {
        consistency: consistencies.length
          ? consistencies.reduce((a, b) => a + b, 0) / consistencies.length
          : null,
        calibration: calibrationFrom(confPairs),
        contaminationDelta:
          anchorAcc != null && freshAcc != null ? Math.max(0, anchorAcc - freshAcc) : null,
        latencyP50: percentile(latencies, 0.5),
        costPer1kItems: costs.length
          ? (costs.reduce((a, b) => a + b, 0) / costs.length) * 1000
          : null,
      };
    }
  }

  // Bradley-Terry matrices for judged duel domains (writing, safety, svg_design).
  const duelMatrices: Record<string, DuelMatrix> = {};
  if (duelBatches.length > 0) {
    for (const domain of DUEL_DOMAINS) {
      const duels = await prisma.pairwiseDuel.findMany({
        where: { runId: { in: duelBatches.map((b) => b.id) }, domain },
        include: {
          modelA: { select: { slug: true } },
          modelB: { select: { slug: true } },
        },
      });
      if (duels.length < 10) continue;
      const slugs = [...new Set(duels.flatMap((d) => [d.modelA.slug, d.modelB.slug]))].sort();
      if (slugs.length < MIN_MODELS) continue;
      const idx = new Map(slugs.map((s, i) => [s, i]));
      const wins = slugs.map(() => slugs.map(() => 0));
      for (const d of duels) {
        const a = idx.get(d.modelA.slug)!;
        const b = idx.get(d.modelB.slug)!;
        if (d.winner === 'a') wins[a]![b]! += 1;
        else if (d.winner === 'b') wins[b]![a]! += 1;
        else {
          wins[a]![b]! += 0.5;
          wins[b]![a]! += 0.5;
        }
      }
      duelMatrices[domain] = { models: slugs, wins };
    }
  }

  const domainsFitted = Object.keys(matrices);
  const duelDomainsFitted = Object.keys(duelMatrices);
  if (domainsFitted.length === 0 && duelDomainsFitted.length === 0) {
    console.warn('[refit] no domain met the minimum matrix size — aborting');
    return null;
  }

  mkdirSync(RUNS_DIR, { recursive: true });
  const stamp = new Date().toISOString().replace(/[:.]/g, '-');
  const inputPath = join(RUNS_DIR, `refit-${stamp}-input.json`);
  const outputPath = join(RUNS_DIR, `refit-${stamp}-output.json`);
  writeFileSync(
    inputPath,
    JSON.stringify(
      { hyperparams: IRT_HYPERPARAMS, domains: matrices, duel_domains: duelMatrices },
      null,
      2,
    ),
  );

  console.log(`[refit] fitting ${domainsFitted.length} domain(s) via ${FIT_SCRIPT}`);
  const py = spawnSync(PYTHON, [FIT_SCRIPT, '--input', inputPath, '--output', outputPath], {
    stdio: 'inherit',
    cwd: REPO_ROOT,
  });
  if (py.status !== 0) throw new Error(`psychometrics fit failed (exit ${py.status})`);

  const fit = JSON.parse(readFileSync(outputPath, 'utf8')) as FitOutput;

  const itemSetHash = createHash('sha256')
    .update(
      domainsFitted.map((d) => matrices[d]!.items.join(',')).join('|') +
        '||' +
        duelDomainsFitted.map((d) => `${d}:${duelMatrices[d]!.models.join(',')}`).join('|'),
    )
    .digest('hex');
  const allSlugs = [
    ...new Set([
      ...domainsFitted.flatMap((d) => matrices[d]!.models),
      ...duelDomainsFitted.flatMap((d) => duelMatrices[d]!.models),
    ]),
  ].sort();

  const run = await prisma.scoreRun.create({
    data: {
      indexVersion: INDEX_VERSION,
      kind: 'index_fit',
      status: 'running',
      itemSetHash,
      modelSet: allSlugs,
      fitDiagnostics: Object.fromEntries([
        ...domainsFitted.map((d) => [d, fit.domains[d]?.diagnostics ?? {}]),
        ...duelDomainsFitted.map((d) => [`duel:${d}`, fit.duel_domains?.[d]?.diagnostics ?? {}]),
      ]) as object,
      notes: `sources=${batchIds.length} eval batches + ${duelBatches.length} duel batches`,
    },
  });

  const perModelDomain = new Map<string, Partial<Record<Domain, ScoreWithCI>>>();

  for (const domain of domainsFitted) {
    const fitted = fit.domains[domain];
    if (!fitted) continue;
    const model = await Promise.all(
      fitted.models.map(async (m) => {
        const dbModel = await prisma.model.findUnique({ where: { slug: m.slug } });
        if (!dbModel) return null;
        const metrics = aux[domain]?.[m.slug];
        const s = domainScore({
          theta: m.theta,
          thetaSe: m.se,
          consistency: metrics?.consistency,
          calibration: metrics?.calibration,
          contaminationDelta: metrics?.contaminationDelta,
        });
        await prisma.score.create({
          data: {
            runId: run.id,
            modelId: dbModel.id,
            domain,
            score: s.score,
            scoreLow: s.scoreLow,
            scoreHigh: s.scoreHigh,
            subMetrics: {
              accuracy_irt: Number((1 / (1 + Math.exp(-m.theta))).toFixed(4)),
              theta: Number(m.theta.toFixed(4)),
              theta_se: Number(m.se.toFixed(4)),
              consistency: metrics?.consistency ?? null,
              calibration: metrics?.calibration ?? null,
              contamination_delta: metrics?.contaminationDelta ?? null,
              latency_p50: metrics?.latencyP50 ?? null,
              cost_per_1k_items: metrics?.costPer1kItems ?? null,
            },
          },
        });
        const acc = perModelDomain.get(m.slug) ?? {};
        acc[domain as Domain] = s;
        perModelDomain.set(m.slug, acc);
        return m.slug;
      }),
    );
    void model;

    // Persist item parameters; discrimination hygiene (§7.4).
    for (const item of fitted.items) {
      const flag =
        item.a < IRT_HYPERPARAMS.minDiscrimination ||
        Math.abs(item.b) > IRT_HYPERPARAMS.maxAbsDifficultyLogits;
      await prisma.evalItem.update({
        where: { id: item.id },
        data: {
          irtA: item.a,
          irtB: item.b,
          ...(flag ? { status: 'flagged_for_retirement' } : {}),
        },
      });
    }
  }

  // Duel domains: Bradley-Terry log-strengths arrive standardized θ-like.
  for (const domain of duelDomainsFitted) {
    const fitted = fit.duel_domains?.[domain];
    if (!fitted) continue;
    for (const m of fitted.models) {
      const dbModel = await prisma.model.findUnique({ where: { slug: m.slug } });
      if (!dbModel) continue;
      const s = domainScore({ theta: m.theta, thetaSe: m.se });
      await prisma.score.create({
        data: {
          runId: run.id,
          modelId: dbModel.id,
          domain,
          score: s.score,
          scoreLow: s.scoreLow,
          scoreHigh: s.scoreHigh,
          subMetrics: {
            accuracy_irt: Number((1 / (1 + Math.exp(-m.theta))).toFixed(4)),
            theta: Number(m.theta.toFixed(4)),
            theta_se: Number(m.se.toFixed(4)),
            method: 'bradley_terry',
          },
        },
      });
      const acc = perModelDomain.get(m.slug) ?? {};
      acc[domain as Domain] = s;
      perModelDomain.set(m.slug, acc);
    }
  }

  for (const [slug, domains] of perModelDomain) {
    const g = globalIndex(domains);
    if (!g) continue;
    const dbModel = await prisma.model.findUnique({ where: { slug } });
    if (!dbModel) continue;
    await prisma.score.create({
      data: {
        runId: run.id,
        modelId: dbModel.id,
        domain: GLOBAL_DOMAIN,
        score: g.score,
        scoreLow: g.scoreLow,
        scoreHigh: g.scoreHigh,
        subMetrics: { domains_covered: Object.keys(domains) },
      },
    });
  }

  await prisma.scoreRun.update({
    where: { id: run.id },
    data: { status: 'complete', completedAt: new Date() },
  });
  console.log(`[refit] index_fit run ${run.id} complete (${domainsFitted.join(', ')})`);
  return { runId: run.id };
}