module CBMReconstructor

using Plots, CSV, DataFrames, Unitful
using Profile, ProfileSVG, BenchmarkTools
using StatsPlots, RDatasets, LinearAlgebra
using Distributed

export Hit, Triplet, Track, 
export load_hits, tr, Dynamic, RK4
export is_neighbour, has_hit, KF_loop
export combination_search, triplets_to_track
export reconstruct_tracks, CA_evolution
export run_reconstruction

struct Hit
	x::Float32
	y::Float32
	z::Int8
end # Hit

struct Triplet
	# 3 hits forming the triplet
	point1::Hit
	point2::Hit
	point3::Hit

	# properties of the triplet
	chi_sq::Float32
	qp::Float32
end

struct Track
	hits::Vector{Hit}
	chi_sq::Float32
end

function load_hits(filename::String)::DataFrame
	# load the CSV into a julia DataFrame
	df = CSV.read(filename, DataFrame)

	select!(df, "x", "y", "z")
	
	# convert the Dataframe columns into numeric types
	df[!, :x] = convert.(Float32, df[:, :x])
	df[!, :y] = convert.(Float32, df[:, :y])
	df[!, :z] = convert.(Int8, df[:, :z])

	# return the converted DataFrame
	return df
	
end # load_hits()

function tr(st)::Float32
    return √(st[3]^2 + st[4]^2 +1)
end # tr(st)

function Dynamic(st)::Vector{Float32}
    c_vacuum = Float32.(299_792_458)
    
    in_ = Vector{Float32}([0.0 for n=1:5])
    
    in_[1] = st[1]
	
    in_[2] = st[2]
	
	# c * (q/p) t_r * (-5 - t_x^2)
    in_[3] = c_vacuum * st[5] * tr(st) * (-5 * st[3]^2) 
	
    # c * (q/p) * t_r * (-5 * t_x * t_y)
	in_[4] = c_vacuum * st[5] * tr(st) * (-5 * st[3] * st[4]) 
	
    in_[5] = 0
    
    return in_
end # Dynamic(st)

function RK4(st, step)::Vector{Float32}
    
    st0 = st
    
    # integrator outputs
    k1::Vector{Float32} = Dynamic(st0)
    k2::Vector{Float32} = Dynamic(st0 + k1 * .5step)
    k3::Vector{Float32} = Dynamic(st0 + k2 * .5step)
    k4::Vector{Float32} = Dynamic(st0 + k3 * step)
    
    return st0 + (k1 + 2k2 + 2k3 + k4) * (step / 6)
    
end # RK4(st, step)

function is_neighbour(A::Triplet, B::Triplet)::Bool
	if A.point2 == B.point1 && A.point3 == B.point2 && (A.qp > 0) == (B.qp > 0)
		return true
	end # if

	return false
end # is_neighbour(A, B)

function has_hit(t::Triplet, h::Hit)::Bool
	if t.point1 == h || t.point2 == h || t.point3 == h
		return true
	else
		return false
	end # if
end # has_hit(t, h)

function KF_loop(r, C, mes, chi_old::Float32, RKstep::Float32)
    
    r_rk = RK4(r, RKstep)
    
    t_x = r_rk[3] / r[3]
    t_y = r_rk[4] / r[4]
    
    F = Float32.([1.0 0.0 RKstep 0.0 0.0
                  0.0 1.0 0.0 RKstep 0.0
                  0.0 0.0 t_x 0.0 0.0
                  0.0 0.0 0.0 t_y 0.0
                  0.0 0.0 0.0 0.0 1.0])
    
    H = Float32.([1.0 0.0 0.0 0.0 0.0 
                  0.0 1.0 0.0 0.0 0.0])
    
    I = Float32.([1.0 0.0 0.0 0.0 0.0
                  0.0 1.0 0.0 0.0 0.0
                  0.0 0.0 1.0 0.0 0.0
                  0.0 0.0 0.0 1.0 0.0
                  0.0 0.0 0.0 0.0 1.0])    
    
    C_pred = F * C * F'
    r_pred = F * r
    
    residual = mes - H * r_pred
    
    W = inv(H * C_pred * H')
    K = C_pred * H' * W
    
    r_corr = r_pred + K * residual
    C_corr = I - K * H * C_pred
    
    chi_new::Float32 = chi_old + residual' * W * residual
    
    return r_corr, C_corr, chi_new
end # KF_loop()

"""
This function accepts two arguments:

1. a DataFrame with three columns (x, y, z) with column types (Float32, Float32, Int8) respectively, containing all the hitson the three detector stations relevant to this part of the search. 

2. a Hit structure, deciding the first Hit to produce triplets for 

The function takes the starting point and calculates all possible triplets, that can be constructed from them while assigning them a χ-square value denoting how likely it is for each triplet to have been produced by one particle
"""
function combination_search(df::DataFrame, start::Hit)::Vector{Triplet}

    r_0 = Float32.([
            0.0; 
            0.0; 
            10 * start.x / start.z; 
            10 * start.y / start.z;
            start.x > 0 ? -1e-19 : 1e-19])
    
    C_0 = Float32.(
        [512.0 0.0 0.0 0.0 0.0
         0.0 512.0 0.0 0.0 0.0
         0.0 0.0 512.0 0.0 0.0
         0.0 0.0 0.0 512.0 0.0
         0.0 0.0 0.0 0.0 512.0])
    
    chi0 = Float32.(0.0)
    
    RKstep = Float32.(start.z / 10.0)
    
    r = Float32.([0.0;0.0;0.0;0.0;0.0])
    
    r, C, chi1 = KF_loop(r_0, C_0, [start.x, start.y], chi0, RKstep)

	# create empty array for the created triplets
	triplets = Vector{Triplet}(undef, 0)
	
	# main parallelization point:
	# take each eiligible starting station
	# and generate triplets for them
	
    for x in eachrow(df[df.z .== start.z + 1, :])

		# save the second point
		point2::Hit = Hit(x...)
		
		# update the RK4 step to represent step of only one station
        step::Float32 = 0.1

		# perform the second iteration of the KF loop 
        r, C, chi2 = KF_loop(r, C, [v for v in values(x[1:2])], chi1, step)

        for y in eachrow(df[df.z .== start.z + 2,:])

			# sace the third point
			point3::Hit = Hit(y...)
			
			# perform the third iteration of the KF loop 
			r, C, chi3 = KF_loop(r, C, [v for v in values(y[1:2])], chi2, step)

			# discard invalid triplets
			if chi3 < 0
				continue
			end
			
			# create the triplet from calculated data
			triplet::Triplet = Triplet(start, point2, point3, chi3, r[5])
			
			# save the triplet
			push!(triplets, triplet)
        end # for
    end # for
	
    return triplets
    
end # combination_search(df, start)

"""
this function expects the triplets to be sorted according to the z axis
"""
function triplets_to_track(triplets::Vector{Triplet})::Track
	hits = [triplets[1].point1, triplets[1].point2, triplets[1].point3]
	chi::Float32 = triplets[1].chi_sq

	# iterate through the rest of the triplets
	for i in 2:length(triplets)
		
		push!(hits, triplets[i].point3)
		chi += triplets[i].chi_sq
	end # for

	return Track(hits, chi)
	
end # triplets_to_track

"""
This is a function that runs the sequential track reconstruction
for a given DataFrame. It gets one positional argument of type DataFrame
"""
function reconstruct_tracks(df::DataFrame)::Vector{Triplet}
	# this is not yet finished, TODO finish this
	
	# create empty array for the created triplets
	triplets = Vector{Triplet}(undef, 0)
	
	# iterate starting stations 1 through 6 (z 3 through 8)
    for i in 3:8
		
        # make needed slices of dataframe (vectorized)
        df_slice = df[(df.z .== i+1) .| (df.z .== i+2), :]
        df_start_points = df[(df.z .== i),:]

		# create the vector of all triplets starting on a station
		triplet_station = Vector{Triplet}(undef, 0)

		# iterate the starting hits on given station
        for x in eachrow(df_start_points)

			# create the starting hit
			start_hit = Hit(x...)

			# iterate the combination search
            result_triplets = combination_search(df_slice, start_hit)

			# append generated triplet list
			append!(triplets, result_triplets)
        end # for
    end # for
	
	return triplets
	
end # reconstruct_tracks

function CA_evolution(triplets::Vector{Triplet})::Vector{Track}

	tracks = Vector{Track}(undef, 0)
	
	# while there still exists a valid triplet, build tracks
	while !isempty(triplets)
		
		track_triplets = [triplets[1]]

		for i in 2:length(triplets)

			if is_neighbour(last(track_triplets), triplets[i])

				push!(track_triplets, triplets[i])
			end # if
		end # for
		push!(tracks, triplets_to_track(track_triplets))

		for h in last(tracks).hits
			filter!(t->!has_hit(t, h), triplets)
		end # for
	end # while

	return tracks
	
end # CA_evolution(triplets)

function run_reconstruction(df::DataFrame)
	sorted_triplets = sort!(reconstruct_tracks(df), by = v -> (v.point1.z, v.chi_sq))
	tracks = CA_evolution(sorted_triplets)
	return tracks
end

end

PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
ProfileSVG = "132c30aa-f267-4189-9183-c8a63c7e05e6"
RDatasets = "ce6b1742-4840-55fa-b093-852dadbb1d8b"
StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"
Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"

[compat]
BenchmarkTools = "~1.3.2"
CSV = "~0.10.9"
DataFrames = "~1.5.0"
PhysicalConstants = "~0.2.3"
Plots = "~1.38.11"
ProfileSVG = "~0.2.1"
RDatasets = "~0.7.7"
StatsPlots = "~0.15.5"
Unitful = "~1.13.1"
"""