import type { Block, Frame, PointDatum, TrailDatum, TripDatum } from "@/types/trajectory";
import { BLOCK_SIZE, MAX_TRAIL_PARTICLES, TRAIL_STEP_STRIDE } from "./constants";

export function buildScatterData(frame: Frame): PointDatum[] {
  const out: PointDatum[] = new Array(frame.length);
  for (let i = 0; i < frame.length; i++) {
    const p = frame[i];
    out[i] = {
      position: [p[0], p[1]],
      beached: p[4] === 1,
      speed: Math.sqrt(p[2] * p[2] + p[3] * p[3]),
      id: i,
    };
  }
  return out;
}

export function buildTrails(
  allBlocks: (Block | null)[],
  currentStep: number,
  numParticles: number,
  releaseSteps: number[] | undefined,
  stepStride: number = TRAIL_STEP_STRIDE
): TrailDatum[] {
  const out: TrailDatum[] = [];
  const maxP = Math.min(numParticles, MAX_TRAIL_PARTICLES);
  for (let pid = 0; pid < maxP; pid++) {
    const releaseStep = releaseSteps?.[pid] ?? 0;
    if (currentStep < releaseStep) continue;
    const path: [number, number][] = [];
    let beached = false;
    for (let s = releaseStep; s <= currentStep; s++) {
      const blockIdx = Math.floor(s / BLOCK_SIZE);
      const block = allBlocks[blockIdx];
      if (!block?.frames?.length) continue;
      const frameIdx = s - block.step_start;
      if (frameIdx < 0 || frameIdx >= block.frames.length) continue;
      const p = block.frames[frameIdx][pid];
      if (!p) continue;
      const isFirst = s === releaseStep;
      const isSubsample = (s - releaseStep) % stepStride === 0;
      const isLast = s === currentStep;
      const isBeached = p[4] === 1;
      if (isFirst || isSubsample || isLast || isBeached) path.push([p[0], p[1]]);
      if (isBeached) {
        beached = true;
        break;
      }
    }
    if (path.length > 0) out.push({ path, beached, id: pid });
  }
  return out;
}

export function buildTripsWithTime(
  allBlocks: (Block | null)[],
  currentStep: number,
  numParticles: number,
  releaseSteps: number[] | undefined,
  stepStride: number = TRAIL_STEP_STRIDE
): TripDatum[] {
  const out: TripDatum[] = [];
  const maxP = Math.min(numParticles, MAX_TRAIL_PARTICLES);
  for (let pid = 0; pid < maxP; pid++) {
    const releaseStep = releaseSteps?.[pid] ?? 0;
    if (currentStep < releaseStep) continue;
    const path: [number, number, number][] = [];
    let beached = false;
    for (let s = releaseStep; s <= currentStep; s++) {
      const blockIdx = Math.floor(s / BLOCK_SIZE);
      const block = allBlocks[blockIdx];
      if (!block?.frames?.length) continue;
      const frameIdx = s - block.step_start;
      if (frameIdx < 0 || frameIdx >= block.frames.length) continue;
      const p = block.frames[frameIdx][pid];
      if (!p) continue;
      const isFirst = s === releaseStep;
      const isSubsample = (s - releaseStep) % stepStride === 0;
      const isLast = s === currentStep;
      const isBeached = p[4] === 1;
      if (isFirst || isSubsample || isLast || isBeached)
        path.push([p[0], p[1], s * 3600]);
      if (isBeached) {
        beached = true;
        break;
      }
    }
    if (path.length > 0) out.push({ path, beached, id: pid });
  }
  return out;
}

/**
 * Build frames from an Arrow block table. Rows must be ordered by step then particle_id.
 * Derives numSteps from the step column so the last block (fewer than block_size steps) is correct.
 */
export function framesFromArrowTable(
  table: import("apache-arrow").Table,
  numParticles: number
): Frame[] {
  const n = numParticles;
  const stepCol = table.getChild("step");
  const lonCol = table.getChild("lon");
  const latCol = table.getChild("lat");
  const uCol = table.getChild("u");
  const vCol = table.getChild("v");
  const beachedCol = table.getChild("beached");
  if (!stepCol || !lonCol || !latCol || !uCol || !vCol || !beachedCol || n <= 0)
    return [];
  const totalRows = table.numRows;
  const minStep = stepCol.get(0) as number;
  const maxStep = stepCol.get(totalRows - 1) as number;
  const numSteps = maxStep - minStep + 1;
  if (numSteps <= 0 || totalRows < numSteps * n) return [];
  const frames: Frame[] = [];
  for (let stepIdx = 0; stepIdx < numSteps; stepIdx++) {
    const frame: number[][] = [];
    for (let pid = 0; pid < n; pid++) {
      const row = stepIdx * n + pid;
      frame.push([
        lonCol.get(row) as number,
        latCol.get(row) as number,
        uCol.get(row) as number,
        vCol.get(row) as number,
        (beachedCol.get(row) as number) ?? 0,
      ]);
    }
    frames.push(frame);
  }
  return frames;
}

export function countBeached(frame: Frame): number {
  let c = 0;
  for (let i = 0; i < frame.length; i++) if (frame[i][4] === 1) c++;
  return c;
}
