import { describe, expect, it } from 'vitest'

import {
  a100_80gb,
  analyzeCluster,
  b300,
  cluster64GPU,
  h100_sxm,
  llama70B,
  llama7B,
  llama31_405B,
  olmo3_32B,
  singleNode8GPU,
  trinityLarge400B,
  type ClusterConfig,
  type TrainingConfig,
} from '../src/lib/trainingClusterModel'

const baselineTraining: TrainingConfig = {
  microBatchSize: 1,
  seqLength: 2048,
  gradAccumSteps: 8,
  precision: 'bf16',
  activationCheckpointing: true,
  optimizer: 'adamw',
}

describe('trainingClusterModel', () => {
  it('fits Llama 2 7B on 8x A100 80GB with TP=8 and derived DP=1', () => {
    const analysis = analyzeCluster(llama7B(), baselineTraining, singleNode8GPU(a100_80gb()), {
      tp: 8,
      pp: 1,
      cp: 1,
      ep: 1,
      distributedOptimizer: false,
      fsdpShardGroupSize: 0,
      zeroStage: 0,
    })

    expect(analysis.feasible).toBe(true)
    expect(analysis.derivedParallelism.dp).toBe(1)
    expect(analysis.memoryBreakdown.totalGB).toBeLessThan(80)
  })

  it('marks Llama 2 70B on 8x A100 80GB as infeasible for unsharded Adam training', () => {
    const analysis = analyzeCluster(llama70B(), baselineTraining, singleNode8GPU(a100_80gb()), {
      tp: 8,
      pp: 1,
      cp: 1,
      ep: 1,
      distributedOptimizer: false,
      fsdpShardGroupSize: 0,
      zeroStage: 0,
    })

    expect(analysis.feasible).toBe(false)
    expect(analysis.infeasibilityReason).toContain('exceeding 80 GB of HBM')
  })

  it('keeps MFU in a realistic range for a balanced 64x H100 dense run', () => {
    const analysis = analyzeCluster(
      llama70B(),
      {
        ...baselineTraining,
        seqLength: 4096,
        gradAccumSteps: 16,
      },
      cluster64GPU(h100_sxm()),
      {
        tp: 4,
        pp: 4,
        cp: 1,
        ep: 1,
        distributedOptimizer: true,
        fsdpShardGroupSize: 0,
        zeroStage: 1,
      },
    )

    expect(analysis.feasible).toBe(true)
    expect(analysis.derivedParallelism.dp).toBe(4)
    expect(analysis.throughput.mfu).toBeGreaterThan(0.3)
    expect(analysis.throughput.mfu).toBeLessThanOrEqual(0.62)
  })

  it('reduces activation memory when CP increases and adds CP communication', () => {
    const withoutCp = analyzeCluster(
      llama70B(),
      {
        ...baselineTraining,
        seqLength: 4096,
      },
      cluster64GPU(h100_sxm()),
      {
        tp: 2,
        pp: 2,
        cp: 1,
        ep: 1,
        distributedOptimizer: true,
        fsdpShardGroupSize: 0,
        zeroStage: 1,
      },
    )
    const withCp = analyzeCluster(
      llama70B(),
      {
        ...baselineTraining,
        seqLength: 4096,
      },
      cluster64GPU(h100_sxm()),
      {
        tp: 2,
        pp: 2,
        cp: 4,
        ep: 1,
        distributedOptimizer: true,
        fsdpShardGroupSize: 0,
        zeroStage: 1,
      },
    )

    expect(withCp.memoryBreakdown.activationsGB).toBeLessThan(withoutCp.memoryBreakdown.activationsGB)
    expect(withCp.communication.cp.totalVolumePerStepGB).toBeGreaterThan(0)
  })

  it('reduces OLMo memory with HSDP shard groups compared with plain DP', () => {
    const cluster = {
      ...cluster64GPU(h100_sxm()),
      numNodes: 128,
      nodesPerRack: 16,
    }
    const plain = analyzeCluster(
      olmo3_32B(),
      {
        microBatchSize: 1,
        seqLength: 8192,
        gradAccumSteps: 1,
        precision: 'bf16',
        activationCheckpointing: true,
        optimizer: 'adamw',
      },
      cluster,
      {
        tp: 1,
        pp: 1,
        cp: 1,
        ep: 1,
        distributedOptimizer: false,
        fsdpShardGroupSize: 0,
        zeroStage: 0,
      },
    )
    const hsdp = analyzeCluster(
      olmo3_32B(),
      {
        microBatchSize: 1,
        seqLength: 8192,
        gradAccumSteps: 1,
        precision: 'bf16',
        activationCheckpointing: true,
        optimizer: 'adamw',
      },
      cluster,
      {
        tp: 1,
        pp: 1,
        cp: 1,
        ep: 1,
        distributedOptimizer: true,
        fsdpShardGroupSize: 256,
        zeroStage: 3,
      },
    )

    expect(hsdp.derivedParallelism.replicaGroups).toBe(4)
    expect(hsdp.memoryBreakdown.totalGB).toBeLessThan(plain.memoryBreakdown.totalGB)
    expect(hsdp.communication.fsdp.totalVolumePerStepGB).toBeGreaterThan(0)
  })

  it('models Trinity as total-parameter-heavy but active-compute-light', () => {
    const analysis = analyzeCluster(
      trinityLarge400B(),
      {
        microBatchSize: 1,
        seqLength: 8192,
        gradAccumSteps: 8,
        precision: 'bf16',
        activationCheckpointing: true,
        optimizer: 'muon',
      },
      trinityCluster(),
      {
        tp: 1,
        pp: 1,
        cp: 1,
        ep: 8,
        distributedOptimizer: true,
        fsdpShardGroupSize: 128,
        zeroStage: 3,
      },
    )

    expect(analysis.totalParams).toBeGreaterThan(300_000_000_000)
    expect(analysis.activeParamsPerToken).toBe(13_000_000_000)
    expect(analysis.communication.ep.totalVolumePerStepGB).toBeGreaterThan(0)
    expect(analysis.communication.ep.usesInterNode).toBe(false)
    expect(new Set(analysis.gpuMap.map((gpu) => gpu.epLane))).toEqual(
      new Set([0, 1, 2, 3, 4, 5, 6, 7]),
    )
  })

  it('derives DP for Llama 3.1 405B from world size and 4D parallelism', () => {
    const analysis = analyzeCluster(
      llama31_405B(),
      {
        microBatchSize: 1,
        seqLength: 8192,
        gradAccumSteps: 16,
        precision: 'bf16',
        activationCheckpointing: true,
        optimizer: 'adamw',
      },
      llama405Cluster(),
      {
        tp: 8,
        pp: 16,
        cp: 1,
        ep: 1,
        distributedOptimizer: true,
        fsdpShardGroupSize: 0,
        zeroStage: 1,
      },
    )

    expect(analysis.derivedParallelism.dp).toBe(128)
    expect(analysis.feasible).toBe(true)
    expect(analysis.communication.tp.totalVolumePerStepGB).toBeGreaterThan(0)
    expect(analysis.communication.pp.totalVolumePerStepGB).toBeGreaterThan(0)
    expect(analysis.communication.fsdp.totalVolumePerStepGB).toBe(0)
  })
})

function llama405Cluster(): ClusterConfig {
  return {
    gpuType: h100_sxm(),
    gpusPerNode: 8,
    numNodes: 2048,
    intraNodeBandwidthGBs: 900,
    interNodeBandwidthGBs: 50,
    nodesPerRack: 16,
    rackLabel: 'rack',
    nodeLabel: 'GPU host',
    podLabel: 'rack',
  }
}

function trinityCluster(): ClusterConfig {
  return {
    gpuType: b300(),
    gpusPerNode: 8,
    numNodes: 256,
    intraNodeBandwidthGBs: 900,
    interNodeBandwidthGBs: 50,
    nodesPerRack: 9,
    rackLabel: 'rack',
    nodeLabel: 'GPU host',
    podLabel: 'rack',
  }
}