import { HARDWARE, hardwareLabel } from "./hardware";

export type Support = "yes" | "no" | "partial" | "unknown";

export type Backend = "cuda" | "rocm" | "metal" | "intel" | "cpu";

export interface Model {
  id: string;
  label: string;
  /** billions of parameters */
  params: number;
  layers: number;
  kvHeads: number;
  headDim: number;
}

export interface Gpu {
  id: string;
  label: string;
  /** GB of VRAM */
  vram: number;
  backend: Backend;
}

export interface Goal {
  id: string;
  label: string;
  /** capability that must be supported for this goal */
  requires?: "peft" | "serializable";
}

export interface Method {
  id: string;
  name: string;
  bits: number[];
  backends: Backend[];
  onTheFly: boolean;
  compile: Support;
  peft: Support;
  serializable: Support;
  /** minutes of setup before you can load */
  setupMinutes: number;
  blurb: string;
  config: string;
  docs: string;
}


/**
 * Every memory configuration in the Hugging Face hardware table, most-owned
 * first — so the cards people actually have sit at the top of the list.
 */
export const GPUS: Gpu[] = HARDWARE.map((h) => ({
  id: h.id,
  label: hardwareLabel(h),
  vram: h.vram,
  backend: h.backend,
}));

export const GOALS: Goal[] = [
  { id: "serve", label: "serve requests" },
  { id: "finetune", label: "fine-tune it", requires: "peft" },
  { id: "publish", label: "publish a checkpoint", requires: "serializable" },
];

export const PRECISIONS = [8, 4, 3, 2] as const;

export const PRECISION_ACCENT: Record<number, string> = {
  8: "#c08532",
  4: "#34785c",
  3: "#f54e00",
  2: "#cf2d56",
};

export const METHODS: Method[] = [
  {
    id: "bitsandbytes",
    name: "bitsandbytes",
    bits: [4, 8],
    backends: ["cuda", "rocm", "metal", "intel", "cpu"],
    onTheFly: true,
    compile: "yes",
    peft: "yes",
    serializable: "yes",
    setupMinutes: 0,
    blurb: "Quantizes as the model loads. No calibration data, no extra pipeline.",
    config: "BitsAndBytesConfig(load_in_4bit=True)",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/bitsandbytes",
  },
  {
    id: "gptqmodel",
    name: "GPTQModel",
    bits: [2, 3, 4, 8],
    backends: ["cuda", "rocm", "metal", "intel", "cpu"],
    onTheFly: false,
    compile: "no",
    peft: "yes",
    serializable: "yes",
    setupMinutes: 25,
    blurb: "Best accuracy available at 4-bit. Costs a calibration pass up front.",
    config: 'GPTQConfig(bits=4, dataset="c4")',
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/gptq",
  },
  {
    id: "awq",
    name: "AWQ",
    bits: [4],
    backends: ["cuda", "rocm", "intel", "cpu"],
    onTheFly: false,
    compile: "unknown",
    peft: "yes",
    serializable: "yes",
    setupMinutes: 20,
    blurb: "Activation-aware and widely deployed. Slightly behind GPTQ on accuracy.",
    config: "AwqConfig(bits=4)",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/awq",
  },
  {
    id: "hqq",
    name: "HQQ",
    bits: [2, 3, 4, 8],
    backends: ["cuda", "intel", "cpu"],
    onTheFly: true,
    compile: "yes",
    peft: "yes",
    serializable: "no",
    setupMinutes: 0,
    blurb: "Fast to apply and needs no dataset, but the result cannot be saved.",
    config: "HqqConfig(nbits=4, group_size=64)",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/hqq",
  },
  {
    id: "compressed-tensors",
    name: "compressed-tensors",
    bits: [4, 8],
    backends: ["cuda", "rocm", "intel", "cpu"],
    onTheFly: false,
    compile: "no",
    peft: "yes",
    serializable: "yes",
    setupMinutes: 0,
    blurb: "Loads checkpoints that are already quantized, and sparse as well.",
    config: "load a compressed-tensors checkpoint",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/compressed_tensors",
  },
  {
    id: "torchao",
    name: "torchao",
    bits: [4, 8],
    backends: ["cuda", "metal", "intel", "cpu"],
    onTheFly: true,
    compile: "unknown",
    peft: "unknown",
    serializable: "partial",
    setupMinutes: 0,
    blurb: "Native to PyTorch. Three of the four capabilities are unverified.",
    config: 'TorchAoConfig("int4_weight_only")',
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/torchao",
  },
  {
    id: "optimum-quanto",
    name: "optimum-quanto",
    bits: [2, 4, 8],
    backends: ["cuda", "metal", "intel", "cpu"],
    onTheFly: true,
    compile: "yes",
    peft: "no",
    serializable: "no",
    setupMinutes: 0,
    blurb: "Portable across backends, including Metal. Cannot save the result.",
    config: "QuantoConfig(weights='int4')",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/quanto",
  },
  {
    id: "higgs",
    name: "HIGGS",
    bits: [2, 4],
    backends: ["cuda"],
    onTheFly: true,
    compile: "yes",
    peft: "no",
    serializable: "yes",
    setupMinutes: 0,
    blurb: "Calibration-free at 2 and 4 bits. No PEFT support.",
    config: "HiggsConfig(bits=4)",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/higgs",
  },
  {
    id: "aqlm",
    name: "AQLM",
    bits: [2],
    backends: ["cuda", "intel", "cpu"],
    onTheFly: false,
    compile: "yes",
    peft: "yes",
    serializable: "yes",
    setupMinutes: 90,
    blurb: "Additive quantization for extreme compression. Long calibration.",
    config: "load an AQLM checkpoint",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/aqlm",
  },
  {
    id: "finegrained-fp8",
    name: "FINEGRAINED_FP8",
    bits: [8],
    backends: ["cuda", "intel"],
    onTheFly: true,
    compile: "no",
    peft: "no",
    serializable: "yes",
    setupMinutes: 0,
    blurb: "Built into Transformers, so there is no extra dependency to pin.",
    config: "FineGrainedFP8Config()",
    docs: "https://huggingface.co/docs/transformers/main/en/quantization/finegrained_fp8",
  },
];

/** Weight bytes, in decimal GB. */
export function weightsGB(model: Model, bits: number): number {
  return (model.params * bits) / 8;
}

/** KV cache for a given context length, in decimal GB (fp16 K and V). */
export function kvCacheGB(model: Model, contextTokens: number): number {
  const bytesPerToken = 2 * model.layers * model.kvHeads * model.headDim * 2;
  return (bytesPerToken * contextTokens) / 1e9;
}

/** Rough allowance for activations and CUDA workspace, in decimal GB. */
export const ACTIVATIONS_GB = 1.5;

export function overheadGB(model: Model, contextTokens: number): number {
  return kvCacheGB(model, contextTokens) + ACTIVATIONS_GB;
}

export function totalGB(model: Model, bits: number, contextTokens: number): number {
  return weightsGB(model, bits) + overheadGB(model, contextTokens);
}

export function methodsFor(bits: number, backend: Backend, goal: Goal): Method[] {
  return METHODS.filter((m) => {
    if (!m.bits.includes(bits)) return false;
    if (!m.backends.includes(backend)) return false;
    if (goal.requires === "peft" && m.peft !== "yes") return false;
    if (goal.requires === "serializable" && m.serializable === "no") return false;
    return true;
  });
}

export function round1(n: number): number {
  return Math.round(n * 10) / 10;
}