'use server';

import { getCapacityPriceTimeline, getDcPriceImpact } from '@/generated/prisma/sql.js';
import { prisma } from '@/lib/db.js';
import { serialize } from '@/lib/superjson.js';
import { validateRegionCode } from '@/lib/utils.js';
import { cacheLife, cacheTag } from 'next/cache';

interface ActionSuccess<T> {
  ok: true;
  data: ReturnType<typeof serialize<T>>;
}

interface ActionError {
  ok: false;
  error: string;
}

type ActionResult<T> = ActionSuccess<T> | ActionError;

// ---------------------------------------------------------------------------
// Multivariate Regression: DC Capacity effect on electricity prices
// controlling for fuel costs, demand, and generation mix
// ---------------------------------------------------------------------------

export interface MultivariateDataPoint {
  region_code: string;
  region_name: string;
  /** Total DC capacity in the region (MW) */
  dc_capacity_mw: number;
  /** Average electricity price ($/MWh) */
  avg_price: number;
  /** Average natural gas price during same period ($/MMBtu) */
  avg_gas_price: number;
  /** Average demand (MW) */
  avg_demand: number;
  /** Gas generation share (fraction 0–1) */
  gas_share: number;
  /** Residual: price unexplained by fuel/demand model */
  residual: number;
}

export interface RegressionCoefficient {
  variable: string;
  coefficient: number;
  /** Standardized coefficient (beta weight) */
  beta: number;
}

export interface MultivariateResult {
  points: MultivariateDataPoint[];
  coefficients: RegressionCoefficient[];
  r2_full: number;
  r2_without_dc: number;
  /** Partial R² attributable to DC capacity */
  r2_partial_dc: number;
  n: number;
}

/**
 * OLS multiple regression using normal equations.
 * X is a matrix of [n x k] predictors (with intercept column prepended).
 * Returns coefficient vector, R², and predicted values.
 */
function olsRegression(
  y: number[],
  X: number[][],
  varNames: string[],
): { coefficients: RegressionCoefficient[]; r2: number; predicted: number[] } | null {
  const n = y.length;
  const k = X[0]!.length;
  if (n <= k) return null;

  // Add intercept column
  const Xa = X.map(row => [1, ...row]);
  const ka = k + 1;

  // X'X
  const XtX: number[][] = Array.from({ length: ka }, () => Array.from<number>({ length: ka }).fill(0));
  for (let i = 0; i < ka; i++) {
    for (let j = 0; j < ka; j++) {
      let sum = 0;
      for (let r = 0; r < n; r++) sum += Xa[r]![i]! * Xa[r]![j]!;
      XtX[i]![j] = sum;
    }
  }

  // X'y
  const Xty: number[] = Array.from<number>({ length: ka }).fill(0);
  for (let i = 0; i < ka; i++) {
    let sum = 0;
    for (let r = 0; r < n; r++) sum += Xa[r]![i]! * y[r]!;
    Xty[i] = sum;
  }

  // Solve via Gauss-Jordan elimination
  const aug: number[][] = XtX.map((row, i) => [...row, Xty[i]!]);
  for (let col = 0; col < ka; col++) {
    // Partial pivoting
    let maxRow = col;
    for (let row = col + 1; row < ka; row++) {
      if (Math.abs(aug[row]![col]!) > Math.abs(aug[maxRow]![col]!)) maxRow = row;
    }
    [aug[col], aug[maxRow]] = [aug[maxRow]!, aug[col]!];
    const pivot = aug[col]![col]!;
    if (Math.abs(pivot) < 1e-12) return null; // singular

    for (let j = col; j <= ka; j++) aug[col]![j]! /= pivot;
    for (let row = 0; row < ka; row++) {
      if (row === col) continue;
      const factor = aug[row]![col]!;
      for (let j = col; j <= ka; j++) aug[row]![j]! -= factor * aug[col]![j]!;
    }
  }

  const beta = aug.map(row => row[ka]!);

  // Predicted and R²
  const predicted = Xa.map(row => row.reduce((sum, val, i) => sum + val * beta[i]!, 0));
  const meanY = y.reduce((s, v) => s + v, 0) / n;
  const ssTot = y.reduce((s, v) => s + (v - meanY) ** 2, 0);
  const ssRes = y.reduce((s, v, i) => s + (v - predicted[i]!) ** 2, 0);
  const r2 = ssTot === 0 ? 0 : 1 - ssRes / ssTot;

  // Compute standardized coefficients (beta weights)
  // beta_std_j = beta_j * (sd_xj / sd_y)
  const sdY = Math.sqrt(ssTot / n);
  const coefficients: RegressionCoefficient[] = varNames.map((name, idx) => {
    const colIdx = idx; // 0-indexed in original X (no intercept)
    const xValues = X.map(row => row[colIdx]!);
    const meanX = xValues.reduce((s, v) => s + v, 0) / n;
    const sdX = Math.sqrt(xValues.reduce((s, v) => s + (v - meanX) ** 2, 0) / n);
    const rawCoeff = beta[idx + 1]!; // +1 to skip intercept
    return {
      variable: name,
      coefficient: rawCoeff,
      beta: sdY > 0 && sdX > 0 ? rawCoeff * (sdX / sdY) : 0,
    };
  });

  return { coefficients, r2, predicted };
}

export async function fetchMultivariateAnalysis(): Promise<ActionResult<MultivariateResult>> {
  'use cache';
  cacheLife('prices');
  cacheTag('multivariate-analysis');

  try {
    // Get all regions with DC data
    const regions = await prisma.gridRegion.findMany({
      select: {
        id: true,
        code: true,
        name: true,
        datacenters: { select: { capacityMw: true } },
      },
    });

    // For each region, compute monthly observations joining price + gas + demand + gen mix
    // This gives us many more data points than just one per region
    const monthlyData = await prisma.$queryRaw<
      Array<{
        region_code: string;
        region_name: string;
        month: Date;
        avg_price: number;
        avg_demand: number;
      }>
    >`
      SELECT
        r.code AS region_code,
        r.name AS region_name,
        date_trunc('month', ep.timestamp) AS month,
        AVG(ep.price_mwh) AS avg_price,
        AVG(ep.demand_mw) AS avg_demand
      FROM electricity_prices ep
      JOIN grid_regions r ON ep.region_id = r.id
      GROUP BY r.code, r.name, date_trunc('month', ep.timestamp)
      HAVING COUNT(*) >= 10
      ORDER BY r.code, month
    `;

    // Get monthly gas prices
    const gasData = await prisma.$queryRaw<Array<{ month: Date; avg_gas: number }>>`
      SELECT
        date_trunc('month', timestamp) AS month,
        AVG(price) AS avg_gas
      FROM commodity_prices
      WHERE commodity = 'natural_gas'
      GROUP BY date_trunc('month', timestamp)
    `;
    const gasByMonth = new Map(gasData.map(g => [g.month.toISOString(), g.avg_gas]));

    // Get gas generation share per region (latest available)
    const genData = await prisma.$queryRaw<Array<{ region_code: string; gas_share: number }>>`
      SELECT
        r.code AS region_code,
        COALESCE(
          SUM(CASE WHEN gm.fuel_type = 'NG' THEN gm.generation_mw ELSE 0 END) /
          NULLIF(SUM(gm.generation_mw), 0),
          0
        ) AS gas_share
      FROM generation_mix gm
      JOIN grid_regions r ON gm.region_id = r.id
      WHERE gm.timestamp >= NOW() - INTERVAL '30 days'
      GROUP BY r.code
    `;
    const gasShareByRegion = new Map(genData.map(g => [g.region_code, Number(g.gas_share)]));

    // DC capacity per region
    const dcCapByRegion = new Map(regions.map(r => [r.code, r.datacenters.reduce((sum, d) => sum + d.capacityMw, 0)]));

    // Build observation matrix: each monthly region observation is a row
    const observations: Array<{
      region_code: string;
      region_name: string;
      price: number;
      gas_price: number;
      demand: number;
      gas_share: number;
      dc_capacity: number;
    }> = [];

    for (const row of monthlyData) {
      const gasPrice = gasByMonth.get(row.month.toISOString());
      if (gasPrice === undefined) continue;

      const gasShare = gasShareByRegion.get(row.region_code) ?? 0;
      const dcCap = dcCapByRegion.get(row.region_code) ?? 0;

      observations.push({
        region_code: row.region_code,
        region_name: row.region_name,
        price: Number(row.avg_price),
        gas_price: Number(gasPrice),
        demand: Number(row.avg_demand),
        gas_share: Number(gasShare),
        dc_capacity: dcCap,
      });
    }

    if (observations.length < 10) {
      return { ok: false, error: 'Insufficient data for multivariate regression' };
    }

    // Full model: price ~ gas_price + demand + gas_share + dc_capacity
    const y = observations.map(o => o.price);
    const X_full = observations.map(o => [o.gas_price, o.demand, o.gas_share, o.dc_capacity]);
    const fullResult = olsRegression(y, X_full, ['gas_price', 'demand', 'gas_share', 'dc_capacity']);

    // Reduced model: price ~ gas_price + demand + gas_share (no DC capacity)
    const X_reduced = observations.map(o => [o.gas_price, o.demand, o.gas_share]);
    const reducedResult = olsRegression(y, X_reduced, ['gas_price', 'demand', 'gas_share']);

    if (!fullResult || !reducedResult) {
      return { ok: false, error: 'Regression failed — singular matrix or insufficient variation' };
    }

    // Partial R² for DC capacity = (R²_full - R²_reduced) / (1 - R²_reduced)
    const r2PartialDc = reducedResult.r2 < 1 ? (fullResult.r2 - reducedResult.r2) / (1 - reducedResult.r2) : 0;

    // Compute residuals from the reduced model (fuel + demand only)
    // These residuals show what's left after accounting for fuel costs and demand
    const residuals = y.map((yi, i) => yi - reducedResult.predicted[i]!);

    // Aggregate to per-region points for the scatter chart
    const regionAgg = new Map<
      string,
      {
        name: string;
        prices: number[];
        residuals: number[];
        gasPrices: number[];
        demands: number[];
        gasShare: number;
        dcCap: number;
      }
    >();

    for (let i = 0; i < observations.length; i++) {
      const obs = observations[i]!;
      const existing = regionAgg.get(obs.region_code) ?? {
        name: obs.region_name,
        prices: [],
        residuals: [],
        gasPrices: [],
        demands: [],
        gasShare: obs.gas_share,
        dcCap: obs.dc_capacity,
      };
      existing.prices.push(obs.price);
      existing.residuals.push(residuals[i]!);
      existing.gasPrices.push(obs.gas_price);
      existing.demands.push(obs.demand);
      regionAgg.set(obs.region_code, existing);
    }

    const points: MultivariateDataPoint[] = Array.from(regionAgg.entries()).map(([code, agg]) => {
      const avg = (arr: number[]) => arr.reduce((s, v) => s + v, 0) / arr.length;
      return {
        region_code: code,
        region_name: agg.name,
        dc_capacity_mw: agg.dcCap,
        avg_price: Number(avg(agg.prices).toFixed(2)),
        avg_gas_price: Number(avg(agg.gasPrices).toFixed(2)),
        avg_demand: Number(avg(agg.demands).toFixed(0)),
        gas_share: Number(agg.gasShare.toFixed(3)),
        residual: Number(avg(agg.residuals).toFixed(2)),
      };
    });

    return {
      ok: true,
      data: serialize({
        points,
        coefficients: fullResult.coefficients,
        r2_full: Number(fullResult.r2.toFixed(4)),
        r2_without_dc: Number(reducedResult.r2.toFixed(4)),
        r2_partial_dc: Number(r2PartialDc.toFixed(4)),
        n: observations.length,
      }),
    };
  } catch (err) {
    return {
      ok: false,
      error: `Failed to run multivariate analysis: ${err instanceof Error ? err.message : String(err)}`,
    };
  }
}

export async function fetchCapacityPriceTimeline(
  regionCode: string,
): Promise<ActionResult<getCapacityPriceTimeline.Result[]>> {
  'use cache';
  cacheLife('prices');
  cacheTag(`capacity-price-timeline-${regionCode}`);

  try {
    if (!validateRegionCode(regionCode)) {
      return { ok: false, error: `Invalid region code: ${regionCode}` };
    }
    const rows = await prisma.$queryRawTyped(getCapacityPriceTimeline(regionCode));
    return { ok: true, data: serialize(rows) };
  } catch (err) {
    return {
      ok: false,
      error: `Failed to fetch capacity/price timeline: ${err instanceof Error ? err.message : String(err)}`,
    };
  }
}

export async function fetchDcPriceImpact(): Promise<ActionResult<getDcPriceImpact.Result[]>> {
  'use cache';
  cacheLife('prices');
  cacheTag('dc-price-impact');

  try {
    const rows = await prisma.$queryRawTyped(getDcPriceImpact());
    return { ok: true, data: serialize(rows) };
  } catch (err) {
    return {
      ok: false,
      error: `Failed to fetch DC price impact: ${err instanceof Error ? err.message : String(err)}`,
    };
  }
}

// ---------------------------------------------------------------------------
// Demand Growth Attribution — DC share of regional load growth
// ---------------------------------------------------------------------------

export interface DemandGrowthRow {
  region_code: string;
  region_name: string;
  /** Average demand in earliest available year (MW) */
  baseline_demand_mw: number;
  /** Average demand in most recent year (MW) */
  current_demand_mw: number;
  /** Total demand growth (MW) */
  demand_growth_mw: number;
  /** Total DC capacity added in this period (MW) */
  dc_capacity_added_mw: number;
  /** Estimated DC share of demand growth (fraction 0–1) */
  dc_share_of_growth: number;
  /** Annualized demand growth rate (%) */
  growth_rate_pct: number;
}

export async function fetchDemandGrowthAttribution(): Promise<ActionResult<DemandGrowthRow[]>> {
  'use cache';
  cacheLife('prices');
  cacheTag('demand-growth-attribution');

  try {
    // Get average demand per region per year
    const yearlyDemand = await prisma.$queryRaw<
      Array<{ region_code: string; region_name: string; yr: number; avg_demand: number }>
    >`
      SELECT
        r.code AS region_code,
        r.name AS region_name,
        EXTRACT(YEAR FROM ep.timestamp)::int AS yr,
        AVG(ep.demand_mw) AS avg_demand
      FROM electricity_prices ep
      JOIN grid_regions r ON ep.region_id = r.id
      GROUP BY r.code, r.name, EXTRACT(YEAR FROM ep.timestamp)
      HAVING COUNT(*) >= 50
      ORDER BY r.code, yr
    `;

    // Get DC capacity by region, grouped by year_opened
    const dcCapacity = await prisma.$queryRaw<
      Array<{ region_code: string; year_opened: number; total_capacity: number }>
    >`
      SELECT
        r.code AS region_code,
        d.year_opened,
        SUM(d.capacity_mw) AS total_capacity
      FROM datacenters d
      JOIN grid_regions r ON d.region_id = r.id
      GROUP BY r.code, d.year_opened
    `;

    // Build per-region results
    const byRegion = new Map<string, typeof yearlyDemand>();
    for (const row of yearlyDemand) {
      const arr = byRegion.get(row.region_code) ?? [];
      arr.push(row);
      byRegion.set(row.region_code, arr);
    }

    const dcByRegion = new Map<string, typeof dcCapacity>();
    for (const row of dcCapacity) {
      const arr = dcByRegion.get(row.region_code) ?? [];
      arr.push(row);
      dcByRegion.set(row.region_code, arr);
    }

    const results: DemandGrowthRow[] = [];

    for (const [code, years] of byRegion.entries()) {
      if (years.length < 2) continue;
      const sorted = years.sort((a, b) => a.yr - b.yr);
      const earliest = sorted[0]!;
      const latest = sorted[sorted.length - 1]!;

      const baselineDemand = Number(earliest.avg_demand);
      const currentDemand = Number(latest.avg_demand);
      const demandGrowth = currentDemand - baselineDemand;
      const yearSpan = latest.yr - earliest.yr;

      // Sum DC capacity added in this period
      const dcEntries = dcByRegion.get(code) ?? [];
      const dcAdded = dcEntries
        .filter(d => d.year_opened >= earliest.yr && d.year_opened <= latest.yr)
        .reduce((sum, d) => sum + Number(d.total_capacity), 0);

      // DC load is typically ~85% of nameplate capacity (high utilization)
      const dcLoadMw = dcAdded * 0.85;
      const dcShareOfGrowth = demandGrowth > 0 ? Math.min(1, dcLoadMw / demandGrowth) : 0;

      const growthRate =
        yearSpan > 0 && baselineDemand > 0 ? ((currentDemand / baselineDemand) ** (1 / yearSpan) - 1) * 100 : 0;

      results.push({
        region_code: code,
        region_name: earliest.region_name,
        baseline_demand_mw: Number(baselineDemand.toFixed(0)),
        current_demand_mw: Number(currentDemand.toFixed(0)),
        demand_growth_mw: Number(demandGrowth.toFixed(0)),
        dc_capacity_added_mw: Number(dcAdded.toFixed(0)),
        dc_share_of_growth: Number(dcShareOfGrowth.toFixed(3)),
        growth_rate_pct: Number(growthRate.toFixed(2)),
      });
    }

    results.sort((a, b) => b.dc_share_of_growth - a.dc_share_of_growth);
    return { ok: true, data: serialize(results) };
  } catch (err) {
    return {
      ok: false,
      error: `Failed to fetch demand growth attribution: ${err instanceof Error ? err.message : String(err)}`,
    };
  }
}