// CarryQuantum.Reference
// F# executable reference backend — the conformance anchor.
// Rust transition/apply_gate must produce identical outputs on the same inputs.
// Run: dotnet run --project QuantumReference.fsproj

module CarryQuantum.Reference

open System
open System.Numerics

// ============================================================
// Complex arithmetic
// ============================================================

type Amp = { Re: float; Im: float }

let ampZero  = { Re = 0.0; Im = 0.0 }
let ampOne   = { Re = 1.0; Im = 0.0 }
let ampI     = { Re = 0.0; Im = 1.0 }

let normSq a  = a.Re * a.Re + a.Im * a.Im
let add a b   = { Re = a.Re + b.Re; Im = a.Im + b.Im }
let mul a b   = { Re = a.Re * b.Re - a.Im * b.Im
                  Im = a.Re * b.Im + a.Im * b.Re }
let scale s a = { Re = s * a.Re; Im = s * a.Im }
let conj a    = { Re = a.Re; Im = -a.Im }
let expI t    = { Re = Math.Cos t; Im = Math.Sin t }

// ============================================================
// State vector  (INV-1: sum of normSq = 1)
// ============================================================

type StateVec = { NumQubits: int; Amps: Amp[] }

let newStateVec n =
    let dim = 1 <<< n
    let amps = Array.create dim ampZero
    amps[0] <- ampOne
    { NumQubits = n; Amps = amps }

// DEF-1: Normalised
let isNormalised (sv: StateVec) =
    let s = sv.Amps |> Array.sumBy normSq
    Math.Abs(s - 1.0) < 1e-9

// ============================================================
// Gates  (INV-1: each gate matrix is unitary)
// ============================================================

let private sqrt2inv = 1.0 / Math.Sqrt 2.0

// Single-qubit gate matrices (2x2, row-major)
let gateMatrix = function
    | "I"  -> [| ampOne;  ampZero; ampZero; ampOne  |]
    | "X"  -> [| ampZero; ampOne;  ampOne;  ampZero |]
    | "Y"  -> [| ampZero; { Re=0.0; Im = -1.0 }; ampI; ampZero |]
    | "Z"  -> [| ampOne;  ampZero; ampZero; { Re = -1.0; Im = 0.0 } |]
    | "H"  -> [| scale sqrt2inv ampOne; scale sqrt2inv ampOne
                 scale sqrt2inv ampOne; scale sqrt2inv { Re = -1.0; Im = 0.0 } |]
    | "S"  -> [| ampOne; ampZero; ampZero; ampI |]
    | "T"  -> [| ampOne; ampZero; ampZero; expI (Math.PI / 4.0) |]
    | name -> failwithf "Unknown gate: %s" name

// Apply a single-qubit gate to qubit `target` in an n-qubit statevec.
// Lifts the 2x2 matrix into the 2^n space via tensor embedding.
let applyGate (gate: string) (target: int) (sv: StateVec) : StateVec =
    let n   = sv.NumQubits
    let dim = 1 <<< n
    let m   = gateMatrix gate
    let out = Array.copy sv.Amps
    for mask in 0 .. (dim >>> 1) - 1 do
        // Insert a 0 at position `target` in `mask`
        let lo = (mask &&& ((1 <<< target) - 1))
        let hi = (mask >>> target) <<< (target + 1)
        let i0 = lo ||| hi          // target bit = 0
        let i1 = i0 ||| (1 <<< target) // target bit = 1
        let a  = sv.Amps[i0]
        let b  = sv.Amps[i1]
        out[i0] <- add (mul m[0] a) (mul m[1] b)
        out[i1] <- add (mul m[2] a) (mul m[3] b)
    { sv with Amps = out }

// Apply CNOT: control = ctrl, target = tgt
let applyCNOT (ctrl: int) (tgt: int) (sv: StateVec) : StateVec =
    let out = Array.copy sv.Amps
    for i in 0 .. sv.Amps.Length - 1 do
        if (i >>> ctrl) &&& 1 = 1 then
            let j = i ^^^ (1 <<< tgt)
            out[i] <- sv.Amps[j]
            out[j] <- sv.Amps[i]
    { sv with Amps = out }

// ============================================================
// Measurement  (Born rule collapse)
// ============================================================

type MeasResult = { Outcome: int; PostState: StateVec }

let measure (qubit: int) (rng: float) (sv: StateVec) : MeasResult =
    // Probability of measuring |1⟩ on `qubit`
    let p1 =
        sv.Amps
        |> Array.indexed
        |> Array.sumBy (fun (i, a) ->
            if (i >>> qubit) &&& 1 = 1 then normSq a else 0.0)
    let outcome = if rng < p1 then 1 else 0
    let norm    = if outcome = 1 then Math.Sqrt p1 else Math.Sqrt (1.0 - p1)
    let collapsed =
        sv.Amps
        |> Array.mapi (fun i a ->
            let bit = (i >>> qubit) &&& 1
            if bit = outcome then scale (1.0 / norm) a else ampZero)
    { Outcome = outcome; PostState = { sv with Amps = collapsed } }

// ============================================================
// FSM  (INV-3, INV-4, INV-5)
// ============================================================

type FSMState =
    | Init | Prepare | Entangle | Compute
    | Measure | Verify | Commit | Halted | CycleLimit

// DEF-4: Allowed transition relation
let allowedTransition (s: FSMState) (s': FSMState) : bool =
    match s, s' with
    | Init,     Prepare   -> true
    | Prepare,  Entangle  -> true
    | Entangle, Compute   -> true
    | Compute,  Measure   -> true
    | Measure,  Verify    -> true
    | Verify,   Commit    -> true
    | Commit,   Prepare   -> true
    | Commit,   Commit    -> true
    | _,        Halted    -> true
    | _,        _         -> false

// DEF-5: Terminal states
let isTerminal = function Halted | CycleLimit -> true | _ -> false

type FSM = { State: FSMState; Cycle: int; MaxCycle: int }

type StepError = | TerminalState | CycleLimitReached | InvalidTransition

// step : FSM → FSMState → Result<FSM, StepError>
// REF-1: only succeeds for AllowedTransition pairs
// REF-2: always fails on terminal states
let step (fsm: FSM) (target: FSMState) : Result<FSM, StepError> =
    if isTerminal fsm.State then
        Error TerminalState                    // INV-5 / REF-2
    elif fsm.Cycle >= fsm.MaxCycle then
        Ok { fsm with State = CycleLimit }     // INV-3 hard ceiling
    elif not (allowedTransition fsm.State target) then
        Error InvalidTransition                // INV-4 / REF-1
    else
        Ok { fsm with State = target; Cycle = fsm.Cycle + 1 }

// ============================================================
// Agent ownership  (INV-7)
// ============================================================

type AgentOwnership = { AgentId: string; Qubits: Set<int> }

let ownershipDisjoint (a: AgentOwnership) (b: AgentOwnership) : bool =
    Set.intersect a.Qubits b.Qubits |> Set.isEmpty

let allAgentsDisjoint (agents: AgentOwnership list) : bool =
    agents |> List.forall (fun a ->
        agents |> List.forall (fun b ->
            a.AgentId = b.AgentId || ownershipDisjoint a b))

// ============================================================
// Conformance harness
// ============================================================

module Conformance =

    let assertNorm (label: string) (sv: StateVec) =
        if not (isNormalised sv) then
            failwithf "NORM VIOLATION after %s: sum_normSq = %f"
                label (sv.Amps |> Array.sumBy normSq)

    let testBellState () =
        let sv = newStateVec 2
        assertNorm "init" sv
        let sv = applyGate "H" 0 sv
        assertNorm "H(0)" sv
        let sv = applyCNOT 0 1 sv
        assertNorm "CNOT(0,1)" sv
        // Bell state: (|00⟩ + |11⟩) / √2
        let expected = 1.0 / Math.Sqrt 2.0
        assert (Math.Abs(normSq sv.Amps[0] |> Math.Sqrt - expected) < 1e-9)
        assert (Math.Abs(normSq sv.Amps[3] |> Math.Sqrt - expected) < 1e-9)
        assert (Math.Abs(normSq sv.Amps[1]) < 1e-9)
        assert (Math.Abs(normSq sv.Amps[2]) < 1e-9)
        printfn "PASS  Bell state normalisation"

    let testFSMTerminalAbsorbing () =
        let fsm = { State = Halted; Cycle = 0; MaxCycle = 100 }
        match step fsm Prepare with
        | Error TerminalState -> printfn "PASS  Halted is absorbing"
        | _ -> failwith "FAIL  Halted allowed illegal transition"

    let testFSMCycleLimitAbsorbing () =
        let fsm = { State = CycleLimit; Cycle = 5; MaxCycle = 100 }
        match step fsm Prepare with
        | Error TerminalState -> printfn "PASS  CycleLimit is absorbing"
        | _ -> failwith "FAIL  CycleLimit allowed illegal transition"

    let testFSMCycleMonotone () =
        let fsm = { State = Init; Cycle = 0; MaxCycle = 10 }
        match step fsm Prepare with
        | Ok fsm' ->
            assert (fsm'.Cycle = fsm.Cycle + 1)
            printfn "PASS  Cycle monotone: %d → %d" fsm.Cycle fsm'.Cycle
        | Error e -> failwithf "FAIL  Unexpected error: %A" e

    let testFSMInvalidTransition () =
        let fsm = { State = Init; Cycle = 0; MaxCycle = 10 }
        match step fsm Compute with  // Init → Compute is not in DAG
        | Error InvalidTransition -> printfn "PASS  Invalid transition rejected"
        | _ -> failwith "FAIL  Invalid transition accepted"

    let testOwnershipDisjoint () =
        let agents = [
            { AgentId = "primary"; Qubits = set [0; 1] }
            { AgentId = "partner"; Qubits = set [2; 3] }
        ]
        assert (allAgentsDisjoint agents)
        printfn "PASS  Agent ownership disjoint"

    let testOwnershipOverlap () =
        let agents = [
            { AgentId = "primary"; Qubits = set [0; 1] }
            { AgentId = "partner"; Qubits = set [1; 2] }  // overlap on qubit 1
        ]
        assert (not (allAgentsDisjoint agents))
        printfn "PASS  Ownership overlap detected"

    let runAll () =
        printfn "=== CarryQuantum Conformance Suite ==="
        testBellState ()
        testFSMTerminalAbsorbing ()
        testFSMCycleLimitAbsorbing ()
        testFSMCycleMonotone ()
        testFSMInvalidTransition ()
        testOwnershipDisjoint ()
        testOwnershipOverlap ()
        printfn "=== All conformance tests passed ==="

[<EntryPoint>]
let main _ =
    Conformance.runAll ()
    0