module advent_of_code.tasks.day_8.Star1

open System
open System.IO

let stream = new StreamReader("tasks/day-8/input.txt")

type Point3D =
    struct
        val X: float
        val Y: float
        val Z: float
        new(s: string) =
            let splitS = s.Split(',') |> Array.map float
            let x, y, z = splitS[0], splitS[1], splitS[2]
            { X = x; Y = y; Z = z }

        member this.GetDistanceFrom(p: Point3D) =
            ((p.X - this.X) ** 2 + (p.Y - this.Y) ** 2 + (p.Z - this.Z) ** 2) |> sqrt

        override this.ToString() =
            $"Point3D X = {this.X} Y = {this.Y} Z = {this.Z}"
    end

type DisjointSet =
    struct
        val Parent: Map<Point3D, Point3D>
        val Size: Map<Point3D, int>
        new(parent: Map<Point3D, Point3D>, size: Map<Point3D, int>) =
            {
                Parent = parent
                Size = size
            }
        new(points: Point3D array) =
            {
                Parent = points |> Array.fold (fun (map: Map<Point3D, Point3D>) (point: Point3D) -> map.Add(point, point)) Map.empty
                Size = points |> Array.fold (fun (map: Map<Point3D, int>) (point: Point3D) -> map.Add(point, 1)) Map.empty
            }
        
        member this.Find (point: Point3D) =
            let p2 = this.Parent[point]
            if (p2 = point) then
                this, point
            else
                let disjoint, p = this.Find p2
                let parent = disjoint.Parent.Add(point, p)
                DisjointSet(parent, this.Size), p
        
        member this.Union (connection: Point3D * Point3D) =
            let p1, p2 = connection.Deconstruct()
            let disjointSet, p1 = this.Find(p1)
            let disjointSet, p2 = disjointSet.Find(p2)
            if p1 = p2 then
                disjointSet
            else
                let parent = disjointSet.Parent.Add(p2, p1)
                let size = disjointSet.Size.Add(p1, disjointSet.Size[p1] + disjointSet.Size[p2]).Add(p2, 0)
                DisjointSet(parent, size)
    end

let connectionFilter (ci: (Point3D * Point3D) * int) =
    let c, i = ci.Deconstruct()
    let p1, p2 = c.Deconstruct()
    i % 2 = 0 && p1.GetDistanceFrom p2 <> 0

let removeIndex (ci: (Point3D * Point3D) * int) =
    let c, _ = ci.Deconstruct()
    c

let sortKey (c: Point3D * Point3D) =
    let p1, p2 = c.Deconstruct()
    p1.GetDistanceFrom p2

let main =
    let points = stream.ReadToEnd().Split('\n') |> Array.map Point3D
    let connections = ((Array.allPairs points points) |> Array.sortBy sortKey |> Array.mapi (fun i el -> el, i) |> Array.filter connectionFilter |> Array.map removeIndex)[.. 999]
    let disjointSet = DisjointSet(points)
    let completeDS = connections |> Array.fold (fun (ds: DisjointSet) (connection: Point3D * Point3D) -> ds.Union(connection)) disjointSet
    let sizes = completeDS.Size.Values |> Seq.sort |> Array.ofSeq
    Array.Reverse sizes
    sizes[..2] |> Array.fold (*) 1