using DataFrames
using LinearAlgebra
using CSV
using DataFramesMeta
# using Plots
# using StatsPlots
using PlotlyJS

begin
    df = CSV.read("compiled.csv", DataFrame)

    df.box = df.bb_length .* df.bb_width .* df.bb_height


    @eachrow! df begin
        @newcol :Ix::Vector{Float64}
        @newcol :Iy::Vector{Float64}
        @newcol :Iz::Vector{Float64}
        @newcol :Ibar::Vector{Float64}

        I = [
            :Ixx :Ixy :Ixz
            :Iyx :Iyy :Iyz
            :Izx :Izy :Izz
        ]

        (:Ix, :Iy, :Iz) = eigvals(I)
        :Ibar = :Ix^2 + :Iy^2 + :Iz^2 |> sqrt

    end



    # Convert material to scalar
    begin
        kv = Dict(reverse.(enumerate(Set(df.material_name))))
        mats = [] |> Vector{Int}
        for material in df.material_name
            push!(mats, kv[material])
        end
        df.material_index = mats
    end

    # Remove outliers
    df = df[df.box.<1e6, :]
    df = df[df.mass.<1000, :]
end

# @df df cornerplot(cols(1:7), compact = true)
features = [:mass, :volume, :density, :area, :box, :Ibar, :material_index]

p1 = plot(df, dimensions = features, kind = "splom", Layout(title = "Raw Data"))

CSV.write("prepped.csv", df)

df.cluster = "Cluster " .* ([1, 3, 2, 1, 2, 1, 1, 3, 1, 3, 2, 3, 1, 1, 2, 2, 1, 3, 1, 3, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 1, 1, 3, 2, 1, 1, 2, 2, 3, 3, 2, 2, 2, 1,] .|> string) # From matlab kmeans idx


p2 = plot(df, dimensions = features, color = :cluster, kind = "splom", Layout(title = "Clustered Data"))



savefig(p1, "prepped.svg", width = 1000, height = 1000)
savefig(p2, "prepped_clustered.svg", width = 1000, height = 1000)

open("./prepped.html", "w") do io
    PlotlyBase.to_html(io, p1.plot)
end
open("./prepped_clustered.html", "w") do io
    PlotlyBase.to_html(io, p2.plot)
end