module Program

open System
open LivingVillage.Kernel
open LivingVillage.Kernel.Sim

let ticksPerSecond = 60L
let secondsPerDay = 86400L
let ticksPerDay = secondsPerDay * ticksPerSecond

let segmentTicks = 900L
let summaryInterval = 5000L

let zeroInput = { MoveX = 0.0f; MoveY = 0.0f }

type Driver =
    { Rng: RngState
      Until: int64
      Input: Input }

let nextSegment (rng: RngState) (until: int64) : Driver =
    let a, r1 = Rng.nextUInt64 rng
    let m, r2 = Rng.nextUInt64 r1
    let diag = 0.70710678f
    let dirX, dirY =
        match int (a &&& 7UL) with
        | 0 -> 1.0f, 0.0f
        | 1 -> diag, diag
        | 2 -> 0.0f, 1.0f
        | 3 -> -diag, diag
        | 4 -> -1.0f, 0.0f
        | 5 -> -diag, -diag
        | 6 -> 0.0f, -1.0f
        | _ -> diag, -diag
    let mag =
        match int (m % 3UL) with
        | 0 -> 0.4f
        | 1 -> 0.7f
        | _ -> 1.0f
    { Rng = r2
      Until = until + segmentTicks
      Input = { MoveX = dirX * mag; MoveY = dirY * mag } }

let actionName (k: NpcActionKind) : string =
    match k with
    | Eat -> "eat"
    | Sleep -> "sleep"
    | Wander -> "wander"
    | Work -> "work"
    | Chat -> "chat"

let npcIndex (id: NpcId) : int =
    match id with NpcId i -> i

let actionDistribution (npcs: Npc[]) : int * int * int * int * int =
    let counts =
        npcs
        |> Array.fold (fun (e, s, w, k, c) npc ->
            match npc.Mind.Action with
            | Eat -> (e + 1, s, w, k, c)
            | Sleep -> (e, s + 1, w, k, c)
            | Wander -> (e, s, w + 1, k, c)
            | Work -> (e, s, w, k + 1, c)
            | Chat -> (e, s, w, k, c + 1)) (0, 0, 0, 0, 0)
    counts

type RunStats =
    { ChatsPerNpc: int64[]
      NonFinite: int64
      OutOfBounds: int64
      World: World }

let runSimulation (verbose: bool) (days: int64) (seed: uint64) (npcCount: int) : RunStats =
    let totalTicks = days * ticksPerDay
    if verbose then
        printfn "living-village headless days=%d seed=%d npcs=%d ticks=%d tickrate=%d" days seed npcCount totalTicks ticksPerSecond
    let mutable world = Sim.initialWorldN seed npcCount
    let mutable driver = { Rng = Rng.ofSeed seed; Until = segmentTicks; Input = zeroInput }
    let maxX = float32 (Sim.mapWidthTiles * Sim.tilePixels - Sim.tilePixels)
    let maxY = float32 (Sim.mapHeightTiles * Sim.tilePixels - Sim.tilePixels)
    let mutable nonFinite = 0L
    let mutable outOfBounds = 0L
    let mutable t = 0L
    let mutable lastAction = (world.Npcs.[0]).Mind.Action
    let mutable lastStart = 0L
    let mutable switchCount = 0L
    let mutable switchDurSum = 0L
    let mutable sleepNight = 0L
    let mutable sleepDay = 0L
    let mutable nightTicks = 0L
    let mutable dayTicks = 0L
    let npcTotal = Array.create npcCount 0L
    let lastActions = Array.create npcCount (world.Npcs.[0]).Mind.Action
    let mutable chatsWindowStart = 0L
    while t < totalTicks do
        let d, input =
            if t < driver.Until then driver, driver.Input
            else
                let d = nextSegment driver.Rng driver.Until
                d, d.Input
        driver <- d
        world <- Sim.step { Input = input } world
        t <- t + 1L
        let nightNow = Sim.isNightTick t
        if nightNow then nightTicks <- nightTicks + 1L else dayTicks <- dayTicks + 1L
        for npc in world.Npcs do
            let idx = npcIndex npc.Id
            if npc.Mind.Action = Chat && lastActions.[idx] <> Chat then
                npcTotal.[idx] <- npcTotal.[idx] + 1L
            lastActions.[idx] <- npc.Mind.Action
        let head = world.Npcs.[0]
        if head.Mind.Action = Sleep then
            if nightNow then sleepNight <- sleepNight + 1L else sleepDay <- sleepDay + 1L
        if head.Mind.Action <> lastAction then
            switchCount <- switchCount + 1L
            switchDurSum <- switchDurSum + (t - lastStart)
            lastStart <- t
            lastAction <- head.Mind.Action
        let p = world.Avatar.Pos
        if Single.IsNaN p.X || Single.IsNaN p.Y || Single.IsInfinity p.X || Single.IsInfinity p.Y then
            nonFinite <- nonFinite + 1L
        if p.X < 0.0f || p.X > maxX || p.Y < 0.0f || p.Y > maxY then outOfBounds <- outOfBounds + 1L
        if verbose && t % summaryInterval = 0L then
            let eat, sleep, wander, work, chat = actionDistribution world.Npcs
            let chatsTotal = Array.sum npcTotal
            let day = float t / float ticksPerDay
            printfn "summary tick=%d day=%.4f eat=%d sleep=%d wander=%d work=%d chat=%d chats_window=%d"
                t day eat sleep wander work chat (chatsTotal - chatsWindowStart)
            chatsWindowStart <- chatsTotal
    if verbose then
        let segments = switchCount + 1L
        let avgActionTicks = float (switchDurSum + (totalTicks - lastStart)) / float segments
        let nightPct = if nightTicks > 0L then 100.0 * float sleepNight / float nightTicks else 0.0
        let dayPct = if dayTicks > 0L then 100.0 * float sleepDay / float dayTicks else 0.0
        let npc0 = world.Npcs.[0]
        let memCount = List.length npc0.Mind.Memory
        let lastValence =
            match List.tryHead npc0.Mind.Memory with
            | Some e -> e.Valence
            | None -> 0.0f
        printfn "npc=%A action_switches=%d avg_action_ticks=%.1f sleep_night_ticks=%d sleep_day_ticks=%d sleep_night_pct=%.2f sleep_day_pct=%.2f mem_count=%d last_valence=%.2f"
            npc0.Id switchCount avgActionTicks sleepNight sleepDay nightPct dayPct memCount lastValence
        let chatsTotal = Array.sum npcTotal
        let minChats = if npcCount > 0 then Array.min npcTotal else 0L
        let maxChats = if npcCount > 0 then Array.max npcTotal else 0L
        let avgChats = if npcCount > 0 then float chatsTotal / float npcCount else 0.0
        printfn "chats_total=%d chats_min=%d chats_max=%d chats_avg=%.2f" chatsTotal minChats maxChats avgChats
        printfn "done tick=%d non-finite=%d out-of-bounds=%d" world.Tick nonFinite outOfBounds
    { ChatsPerNpc = npcTotal
      NonFinite = nonFinite
      OutOfBounds = outOfBounds
      World = world }

let dumpRelations (world: World) : unit =
    let m = Sim.relationMatrix world
    let n = Array2D.length1 m
    let ci = System.Globalization.CultureInfo.InvariantCulture
    printfn "relations tick=%d npcs=%d half_life_ticks=%d threshold=%.2f" world.Tick n Sim.relationHalfLifeTicks Sim.relationThreshold
    for i in 0 .. n - 1 do
        let row = Array.init n (fun j -> m.[i, j].ToString("0.000", ci))
        printfn "rel[%d] %s" i (String.concat " " row)
    let counts = Sim.relationCounts m
    let cells = Array.mapi (fun i c -> $"{i}:{c}") counts
    printfn "rel_counts %s" (String.concat " " cells)
    printfn "rel_total=%d rel_max=%d rel_min=%d" (Array.sum counts) (Array.max counts) (Array.min counts)

let gini (xs: int64[]) : float =
    let n = float xs.Length
    let s = float (Array.sum xs)
    if n <= 1.0 || s <= 0.0 then 0.0
    else
        let mutable num = 0.0
        for xi in xs do
            for xj in xs do
                num <- num + abs (float xi - float xj)
        num / (2.0 * n * n * (s / n))

let runBatch (worlds: int) (days: int64) : int =
    printfn "batch worlds=%d days=%d half_life_ticks=%d rel_threshold=%.2f" worlds days Sim.relationHalfLifeTicks Sim.relationThreshold
    printfn "checks: (a) chats max/avg > 2.0 (b) rel_counts max > 2x min (c) no NaN"
    let mutable passed = 0
    let mutable failed = 0
    let ratios = ResizeArray<float> ()
    let ginis = ResizeArray<float> ()
    let sw = System.Diagnostics.Stopwatch.StartNew()
    for k in 0 .. worlds - 1 do
        let seed = 42UL + uint64 k
        let stats = runSimulation false days seed Sim.npcCount
        let m = Sim.relationMatrix stats.World
        let n = Array2D.length1 m
        let mutable hasNonFinite = stats.NonFinite > 0L
        for i in 0 .. n - 1 do
            for j in 0 .. n - 1 do
                if Single.IsNaN m.[i, j] || Single.IsInfinity m.[i, j] then hasNonFinite <- true
        let chats = stats.ChatsPerNpc
        let chatsTotal = Array.sum chats
        let chatsAvg = float chatsTotal / float chats.Length
        let chatsMax = if chats.Length > 0 then Array.max chats else 0L
        let ratio = if chatsAvg > 0.0 then float chatsMax / chatsAvg else 0.0
        let g = gini chats
        let counts = Sim.relationCounts m
        let cntMax = if counts.Length > 0 then Array.max counts else 0
        let cntMin = if counts.Length > 0 then Array.min counts else 0
        let checkA = ratio > 2.0
        let checkB = float cntMax > 2.0 * float cntMin
        let checkC = not hasNonFinite
        let ok = checkA && checkB && checkC
        if ok then passed <- passed + 1 else failed <- failed + 1
        ratios.Add ratio
        ginis.Add g
        printfn "world=%d seed=%d chats_max=%d chats_avg=%.2f ratio=%.3f gini=%.3f relcnt_max=%d relcnt_min=%d nonfinite=%d check_a=%s check_b=%s check_c=%s %s"
            k seed chatsMax chatsAvg ratio g cntMax cntMin stats.NonFinite
            (if checkA then "PASS" else "FAIL")
            (if checkB then "PASS" else "FAIL")
            (if checkC then "PASS" else "FAIL")
            (if ok then "OK" else "BAD")
    sw.Stop()
    let allPass = failed = 0
    printfn "batch_summary worlds=%d passed=%d failed=%d ratio_min=%.3f ratio_max=%.3f gini_min=%.3f gini_max=%.3f elapsed_s=%.1f"
        worlds passed failed (Seq.min ratios) (Seq.max ratios) (Seq.min ginis) (Seq.max ginis) sw.Elapsed.TotalSeconds
    printfn "M3_ACCEPTANCE=%s" (if allPass then "PASS" else "FAIL")
    if allPass then 0 else 1

[<EntryPoint>]
let main argv =
    let rec parse
        (i: int)
        (days: int64 option)
        (seed: uint64 option)
        (npc: int option)
        (dump: bool)
        (batch: (int * int64) option)
        : Result<int64 * uint64 * int * bool * (int * int64) option, string> =
        if i >= argv.Length then
            match batch with
            | Some(k, d) -> Ok(defaultArg days 1L, defaultArg seed 42UL, defaultArg npc 30, dump, Some(k, d))
            | None ->
                match days, seed, npc with
                | Some d, Some s, Some n -> Ok(d, s, n, dump, None)
                | _ -> Error "missing --days/--seed"
        else
            match argv.[i] with
            | "--days" when i + 1 < argv.Length ->
                (match Int64.TryParse argv.[i + 1] with
                 | true, d when d > 0L -> parse (i + 2) (Some d) seed npc dump batch
                 | _ -> Error $"invalid --days '{argv.[i + 1]}'")
            | "--seed" when i + 1 < argv.Length ->
                (match UInt64.TryParse argv.[i + 1] with
                 | true, s -> parse (i + 2) days (Some s) npc dump batch
                 | _ -> Error $"invalid --seed '{argv.[i + 1]}'")
            | "--npc" when i + 1 < argv.Length ->
                (match Int32.TryParse argv.[i + 1] with
                 | true, n when n > 0 && n <= 1000 -> parse (i + 2) days seed (Some n) dump batch
                 | _ -> Error $"invalid --npc '{argv.[i + 1]}'")
            | "--dump-relations" -> parse (i + 1) days seed npc true batch
            | "--batch" when i + 2 < argv.Length ->
                (match Int32.TryParse argv.[i + 1], Int64.TryParse argv.[i + 2] with
                 | (true, k), (true, d) when k > 0 && d > 0L -> parse (i + 3) days seed npc dump (Some(k, d))
                 | _ -> Error $"invalid --batch '{argv.[i + 1]} {argv.[i + 2]}'")
            | other -> Error $"unknown argument '{other}'"

    match parse 0 None None (Some 30) false None with
    | Ok(days, seed, npc, dump, batch) ->
        match batch with
        | Some(k, d) -> runBatch k d
        | None ->
            let stats = runSimulation true days seed npc
            if dump then dumpRelations stats.World
            if stats.NonFinite > 0L || stats.OutOfBounds > 0L then 1 else 0
    | Error msg ->
        eprintfn $"headless: {msg}"
        eprintfn "usage: dotnet run -c Release --project src/LivingVillage.Headless -- --days N --seed S [--npc N] [--dump-relations] | --batch K D"
        2