import type { IRow, IViewField } from '../../../interfaces';
import { getMeaAggKey } from '../../../utils';
import { EXPLAIN_COUNT_KEY, isMarkRow } from '../queries';
import { fitRidge, looCvR2, mulberry32, permutationPValue, predictRidge, selectLambda, shuffled, standardize } from '../stats/ridge';
import { GROUP_COLOR_SCALE, GROUP_CONTEXT, GROUP_SELECTED, LEGEND_BOTTOM } from '../vizTheme';
import type { IExplainContext, IExplanationStrength } from '../types';

/**
 * Shared base-model / explain-model framework for the model-based
 * explainers (contributing dimension & contributing measure), following the
 * Tableau Explain Data approach: hold out the selected mark, compare
 * cross-validated fit and held-out prediction with vs without the candidate
 * feature, and calibrate against permutations of the candidate.
 *
 * SCORING-FROZEN for the executor agent (see docs/explain-data-worklog.md §1).
 */

export const MIN_MARKS = 8;
export const PERMUTATIONS = 30;
export const PERMUTATION_P_THRESHOLD = 0.05;
export const MIN_SCORE = 0.15;

/** target aggregations we can reconstruct from additive sum/count queries */
const DERIVABLE_AGGS = new Set(['sum', 'mean', 'count']);

export function targetMeasure(ctx: IExplainContext): IViewField | null {
    return ctx.viewMeasures.find((m) => DERIVABLE_AGGS.has(m.aggName ?? 'sum')) ?? null;
}

export interface IMarksTable {
    /** one entry per mark, stable order */
    markKeys: string[];
    /** index of the selected mark within markKeys */
    selectedIndex: number;
    /** target measure value per mark (respecting its aggregation) */
    y: number[];
    /** per-mark row counts */
    counts: number[];
    /** raw grouped rows bucketed per mark (for feature construction) */
    rowsByMark: IRow[][];
}

export function markKeyOf(row: IRow, ctx: IExplainContext): string {
    return ctx.viewDimensions.map((d) => String(row[d.fid])).join(' ');
}

/**
 * Bucket view-LOD (or finer) grouped rows into marks and derive the target
 * measure per mark from additive sums/counts. Returns null when the marks
 * table is unusable for modeling (too few marks, selected mark missing…).
 */
export function buildMarksTable(ctx: IExplainContext, rows: IRow[], target: IViewField): IMarksTable | null {
    const sumKey = getMeaAggKey(target.fid, 'sum');
    const byMark = new Map<string, { rows: IRow[]; sum: number; count: number; isSelected: boolean }>();
    for (const row of rows) {
        const key = markKeyOf(row, ctx);
        let bucket = byMark.get(key);
        if (!bucket) {
            bucket = { rows: [], sum: 0, count: 0, isSelected: isMarkRow(row, ctx) };
            byMark.set(key, bucket);
        }
        bucket.rows.push(row);
        bucket.sum += Number(row[sumKey]) || 0;
        bucket.count += Number(row[EXPLAIN_COUNT_KEY]) || 0;
    }
    const entries = [...byMark.entries()];
    if (entries.length < MIN_MARKS) return null;
    const selectedIndex = entries.findIndex(([, b]) => b.isSelected);
    if (selectedIndex < 0) return null;
    const agg = target.aggName ?? 'sum';
    const y = entries.map(([, b]) => {
        if (agg === 'count') return b.count;
        if (agg === 'mean') return b.count > 0 ? b.sum / b.count : 0;
        return b.sum;
    });
    return {
        markKeys: entries.map(([k]) => k),
        selectedIndex,
        y,
        counts: entries.map(([, b]) => b.count),
        rowsByMark: entries.map(([, b]) => b.rows),
    };
}

/**
 * One-hot base features from the view dimensions (drop-first per dimension).
 * When the base would be near-saturated (p > n/2 — e.g. a single view
 * dimension, where dummies memorize each mark), fall back to intercept-only:
 * cross-validated fit then honestly reflects "no generalizable structure".
 */
export function baseFeatures(ctx: IExplainContext, table: IMarksTable): number[][] {
    const n = table.markKeys.length;
    const cols: number[][] = [];
    for (const dim of ctx.viewDimensions) {
        const values = table.rowsByMark.map((rows) => String(rows[0]?.[dim.fid]));
        const distinct = [...new Set(values)];
        for (const v of distinct.slice(1)) {
            cols.push(values.map((x) => (x === v ? 1 : 0)));
        }
    }
    if (cols.length > n / 2) return table.markKeys.map(() => []);
    // transpose to row-major
    return table.markKeys.map((_, r) => cols.map((col) => col[r]));
}

export interface ICandidateScore {
    score: number;
    fitGain: number;
    predGain: number;
    pValue: number;
    cvBase: number;
    cvExplain: number;
}

const clamp01 = (v: number) => Math.min(1, Math.max(0, v));

/** deterministic string hash for permutation seeding */
export function seedFromString(s: string): number {
    let h = 5381;
    for (let i = 0; i < s.length; i++) h = ((h << 5) + h + s.charCodeAt(i)) >>> 0;
    return h;
}

/**
 * Score a candidate feature block against the base model:
 *   fitGain  — LOO-CV R² improvement on the non-selected marks
 *   predGain — held-out (selected mark) absolute-error improvement, in
 *              units of the training target's standard deviation
 *   pValue   — permutation calibration of the combined score
 */
export function scoreCandidate(baseX: number[][], candX: number[][], y: number[], selectedIndex: number, seed: number): ICandidateScore | null {
    const n = y.length;
    if (n - 1 < MIN_MARKS - 1 || candX.length !== n || baseX.length !== n) return null;
    const trainIdx = y.map((_, i) => i).filter((i) => i !== selectedIndex);
    const pExplain = (baseX[0]?.length ?? 0) + (candX[0]?.length ?? 0);
    if (pExplain >= trainIdx.length - 2) return null; // not enough marks to support the feature count

    const explainX = y.map((_, i) => [...baseX[i], ...candX[i]]);
    const yTrain = trainIdx.map((i) => y[i]);
    const sdY = Math.sqrt(yTrain.reduce((acc, v, _, arr) => acc + (v - arr.reduce((a, b) => a + b, 0) / arr.length) ** 2, 0) / yTrain.length) || 1;

    const evaluate = (bX: number[][], eX: number[][]): { fitGain: number; predGain: number; score: number; cvBase: number; cvExplain: number } => {
        const bTrain = trainIdx.map((i) => bX[i]);
        const eTrain = trainIdx.map((i) => eX[i]);
        const { lambda: lb, cvR2: cvBase } = selectLambda(bTrain, yTrain);
        const { lambda: le, cvR2: cvExplain } = selectLambda(eTrain, yTrain);
        const fitGain = cvExplain - cvBase;
        const baseModel = fitRidge(bTrain, yTrain, lb);
        const explainModel = fitRidge(eTrain, yTrain, le);
        const errBase = Math.abs(predictRidge(baseModel, [bX[selectedIndex]])[0] - y[selectedIndex]);
        const errExplain = Math.abs(predictRidge(explainModel, [eX[selectedIndex]])[0] - y[selectedIndex]);
        const predGain = (errBase - errExplain) / sdY;
        const score = 0.5 * clamp01(fitGain) + 0.5 * clamp01(predGain);
        return { fitGain, predGain, score, cvBase, cvExplain };
    };

    const real = evaluate(baseX, explainX);
    if (real.score <= 0) return { ...real, pValue: 1 };

    const rand = mulberry32(seed);
    const permScores: number[] = [];
    for (let b = 0; b < PERMUTATIONS; b++) {
        const perm = shuffled(
            y.map((_, i) => i),
            rand
        );
        const permExplainX = y.map((_, i) => [...baseX[i], ...candX[perm[i]]]);
        permScores.push(evaluate(baseX, permExplainX).score);
    }
    return { ...real, pValue: permutationPValue(real.score, permScores) };
}

export function modelStrength(score: number): IExplanationStrength {
    if (score >= 0.5) return 'strong';
    if (score >= 0.25) return 'moderate';
    return 'weak';
}

/** standardized column + its standardized square (second-order term) */
export function secondOrderFeatures(values: number[]): number[][] {
    const z = standardize(values);
    const z2 = standardize(z.map((v) => v * v));
    return values.map((_, i) => [z[i], z2[i]]);
}

/** marks scatter evidence: candidate feature vs target, gray context + accent selected */
export function scatterSpec(xValues: number[], y: number[], selectedIndex: number, xTitle: string, yTitle: string): Record<string, unknown> {
    return {
        width: 'container',
        height: 220,
        data: {
            values: xValues.map((x, i) => ({ x, y: y[i], group: i === selectedIndex ? GROUP_SELECTED : GROUP_CONTEXT })),
        },
        mark: { type: 'point', filled: true },
        encoding: {
            x: { field: 'x', type: 'quantitative', title: xTitle },
            y: { field: 'y', type: 'quantitative', title: yTitle },
            color: { field: 'group', type: 'nominal', scale: GROUP_COLOR_SCALE, legend: LEGEND_BOTTOM },
            size: { condition: { test: `datum.group === '${GROUP_SELECTED}'`, value: 180 }, value: 55 },
            opacity: { condition: { test: `datum.group === '${GROUP_SELECTED}'`, value: 1 }, value: 0.75 },
            order: { condition: { test: `datum.group === '${GROUP_SELECTED}'`, value: 1 }, value: 0 },
        },
    };
}