package kalman

import "math"

const (
	weightPosition = 1.0 / 20
	weightVelocity = 1.0 / 160
)

type Filter struct {
	X [8]float64
	P [8][8]float64
}

func NewFilter() *Filter {
	kf := &Filter{}
	for i := 0; i < 8; i++ {
		kf.P[i][i] = 10.0
	}
	return kf
}

func (kf *Filter) Predict() {
	motionMat := [8][8]float64{
		{1, 0, 0, 0, 1, 0, 0, 0},
		{0, 1, 0, 0, 0, 1, 0, 0},
		{0, 0, 1, 0, 0, 0, 1, 0},
		{0, 0, 0, 1, 0, 0, 0, 1},
		{0, 0, 0, 0, 1, 0, 0, 0},
		{0, 0, 0, 0, 0, 1, 0, 0},
		{0, 0, 0, 0, 0, 0, 1, 0},
		{0, 0, 0, 0, 0, 0, 0, 1},
	}

	stdPos := [4]float64{
		weightPosition * kf.X[2],
		weightPosition * kf.X[3],
		weightPosition * kf.X[2],
		weightPosition * kf.X[3],
	}
	stdVel := [4]float64{
		weightVelocity * kf.X[2],
		weightVelocity * kf.X[3],
		weightVelocity * kf.X[2],
		weightVelocity * kf.X[3],
	}

	var Q [8][8]float64
	full := [8]float64{stdPos[0], stdPos[1], stdPos[2], stdPos[3], stdVel[0], stdVel[1], stdVel[2], stdVel[3]}
	for i := 0; i < 8; i++ {
		Q[i][i] = full[i] * full[i]
	}

	kf.X = mul8x8_8x1(motionMat, kf.X)
	kf.P = add8x8(mul8x8_8x8(mul8x8_8x8(motionMat, kf.P), transpose8(motionMat)), Q)
}

func (kf *Filter) Update(z [4]float64) {
	updateMat := [4][8]float64{
		{1, 0, 0, 0, 0, 0, 0, 0},
		{0, 1, 0, 0, 0, 0, 0, 0},
		{0, 0, 1, 0, 0, 0, 0, 0},
		{0, 0, 0, 1, 0, 0, 0, 0},
	}

	Rdiag := [4]float64{
		weightPosition * z[2],
		weightPosition * z[3],
		weightPosition * z[2],
		weightPosition * z[3],
	}
	var R [4][4]float64
	for i := 0; i < 4; i++ {
		R[i][i] = Rdiag[i] * Rdiag[i]
	}

	H := updateMat

	HP := mul4x8_8x8(H, kf.P)

	Ht := transpose4x8(H)
	HPHt := add4x4(mul4x8_8x4(HP, Ht), R)

	sinv := inv4x4(HPHt)

	PHt := mul8x8_8x4(kf.P, Ht)

	K := mul8x4_4x4(PHt, sinv)

	y := [4]float64{
		z[0] - dot8(H[0][:], kf.X[:]),
		z[1] - dot8(H[1][:], kf.X[:]),
		z[2] - dot8(H[2][:], kf.X[:]),
		z[3] - dot8(H[3][:], kf.X[:]),
	}

	var Ky [8]float64
	for i := 0; i < 8; i++ {
		for j := 0; j < 4; j++ {
			Ky[i] += K[i][j] * y[j]
		}
	}
	for i := 0; i < 8; i++ {
		kf.X[i] += Ky[i]
	}

	var KH [8][8]float64
	for i := 0; i < 8; i++ {
		for k := 0; k < 4; k++ {
			for j := 0; j < 8; j++ {
				KH[i][j] += K[i][k] * H[k][j]
			}
		}
	}
	var IKH [8][8]float64
	for i := 0; i < 8; i++ {
		IKH[i][i] = 1.0
		for j := 0; j < 8; j++ {
			IKH[i][j] -= KH[i][j]
		}
	}

	IKHP := mul8x8_8x8(IKH, kf.P)
	IKHKt := mul8x8_8x8(IKHP, transpose8(IKH))

	KR := mul8x4_4x4(K, R)
	KRKt := mul8x4_4x8(KR, transpose8x4(K))

	kf.P = add8x8(IKHKt, KRKt)
}

type Tracker struct {
	ID           int
	Filter       *Filter
	TimeSinceUpd int
	Hits         int
	HitStreak    int
	Age          int
}

var nextID int

func NewTracker(bbox [4]float64) *Tracker {
	nextID++
	x := (bbox[0] + bbox[2]) / 2
	y := (bbox[1] + bbox[3]) / 2
	w := bbox[2] - bbox[0]
	h := bbox[3] - bbox[1]

	trk := &Tracker{
		ID:     nextID,
		Filter: NewFilter(),
		Hits:   1,
	}
	trk.Filter.X = [8]float64{x, y, w, h, 0, 0, 0, 0}
	return trk
}

func (t *Tracker) Predict() {
	if t.Filter.X[6]+t.Filter.X[2] <= 0 {
		t.Filter.X[6] = 0
	}
	t.Filter.Predict()
	t.Age++
	t.TimeSinceUpd++
}

func (t *Tracker) Update(bbox [4]float64) {
	t.TimeSinceUpd = 0
	t.Hits++
	t.HitStreak++

	x := (bbox[0] + bbox[2]) / 2
	y := (bbox[1] + bbox[3]) / 2
	w := bbox[2] - bbox[0]
	h := bbox[3] - bbox[1]

	t.Filter.Update([4]float64{x, y, w, h})
}

func (t *Tracker) GetState() [4]float64 {
	xx := t.Filter.X
	x, y, w, h := xx[0], xx[1], xx[2], xx[3]
	return [4]float64{
		x - w/2,
		y - h/2,
		x + w/2,
		y + h/2,
	}
}

func (t *Tracker) GetCX() float64 {
	return t.Filter.X[0]
}

func mul8x8_8x1(m [8][8]float64, v [8]float64) [8]float64 {
	var r [8]float64
	for i := 0; i < 8; i++ {
		for j := 0; j < 8; j++ {
			r[i] += m[i][j] * v[j]
		}
	}
	return r
}

func mul8x8_8x8(a, b [8][8]float64) [8][8]float64 {
	var r [8][8]float64
	for i := 0; i < 8; i++ {
		for k := 0; k < 8; k++ {
			aik := a[i][k]
			for j := 0; j < 8; j++ {
				r[i][j] += aik * b[k][j]
			}
		}
	}
	return r
}

func mul8x8_8x4(a [8][8]float64, b [8][4]float64) [8][4]float64 {
	var r [8][4]float64
	for i := 0; i < 8; i++ {
		for k := 0; k < 8; k++ {
			aik := a[i][k]
			for j := 0; j < 4; j++ {
				r[i][j] += aik * b[k][j]
			}
		}
	}
	return r
}

func mul8x4_4x4(a [8][4]float64, b [4][4]float64) [8][4]float64 {
	var r [8][4]float64
	for i := 0; i < 8; i++ {
		for k := 0; k < 4; k++ {
			aik := a[i][k]
			for j := 0; j < 4; j++ {
				r[i][j] += aik * b[k][j]
			}
		}
	}
	return r
}

func mul8x4_4x8(a [8][4]float64, b [4][8]float64) [8][8]float64 {
	var r [8][8]float64
	for i := 0; i < 8; i++ {
		for k := 0; k < 4; k++ {
			aik := a[i][k]
			for j := 0; j < 8; j++ {
				r[i][j] += aik * b[k][j]
			}
		}
	}
	return r
}

func mul4x8_8x8(a [4][8]float64, b [8][8]float64) [4][8]float64 {
	var r [4][8]float64
	for i := 0; i < 4; i++ {
		for k := 0; k < 8; k++ {
			aik := a[i][k]
			for j := 0; j < 8; j++ {
				r[i][j] += aik * b[k][j]
			}
		}
	}
	return r
}

func mul4x8_8x4(a [4][8]float64, b [8][4]float64) [4][4]float64 {
	var r [4][4]float64
	for i := 0; i < 4; i++ {
		for k := 0; k < 8; k++ {
			aik := a[i][k]
			for j := 0; j < 4; j++ {
				r[i][j] += aik * b[k][j]
			}
		}
	}
	return r
}

func add8x8(a, b [8][8]float64) [8][8]float64 {
	var r [8][8]float64
	for i := 0; i < 8; i++ {
		for j := 0; j < 8; j++ {
			r[i][j] = a[i][j] + b[i][j]
		}
	}
	return r
}

func add4x4(a, b [4][4]float64) [4][4]float64 {
	var r [4][4]float64
	for i := 0; i < 4; i++ {
		for j := 0; j < 4; j++ {
			r[i][j] = a[i][j] + b[i][j]
		}
	}
	return r
}

func transpose8(a [8][8]float64) [8][8]float64 {
	var r [8][8]float64
	for i := 0; i < 8; i++ {
		for j := 0; j < 8; j++ {
			r[i][j] = a[j][i]
		}
	}
	return r
}

func transpose4x8(a [4][8]float64) [8][4]float64 {
	var r [8][4]float64
	for i := 0; i < 4; i++ {
		for j := 0; j < 8; j++ {
			r[j][i] = a[i][j]
		}
	}
	return r
}

func transpose8x4(a [8][4]float64) [4][8]float64 {
	var r [4][8]float64
	for i := 0; i < 8; i++ {
		for j := 0; j < 4; j++ {
			r[j][i] = a[i][j]
		}
	}
	return r
}

func inv4x4(a [4][4]float64) [4][4]float64 {
	var inv [4][4]float64

	m := [4][4]float64{
		{a[0][0], a[0][1], a[0][2], a[0][3]},
		{a[1][0], a[1][1], a[1][2], a[1][3]},
		{a[2][0], a[2][1], a[2][2], a[2][3]},
		{a[3][0], a[3][1], a[3][2], a[3][3]},
	}

	col := [4]int{0, 1, 2, 3}
	row := [4]int{0, 1, 2, 3}

	for i := 0; i < 4; i++ {
		maxVal := math.Abs(m[row[i]][col[i]])
		pi, pj := i, i
		for r := i; r < 4; r++ {
			for c := i; c < 4; c++ {
				v := math.Abs(m[row[r]][col[c]])
				if v > maxVal {
					maxVal = v
					pi, pj = r, c
				}
			}
		}
		row[i], row[pi] = row[pi], row[i]
		col[i], col[pj] = col[pj], col[i]

		pivot := m[row[i]][col[i]]
		if math.Abs(pivot) < 1e-12 {
			pivot = 1e-12
		}
		m[row[i]][col[i]] = 1.0
		for j := 0; j < 4; j++ {
			m[row[i]][j] /= pivot
		}

		for r := 0; r < 4; r++ {
			if r != i {
				factor := m[row[r]][col[i]]
				m[row[r]][col[i]] = 0
				for j := 0; j < 4; j++ {
					m[row[r]][j] -= factor * m[row[i]][j]
				}
			}
		}
	}

	for i := 0; i < 4; i++ {
		for j := 0; j < 4; j++ {
			inv[col[i]][row[j]] = m[row[i]][j]
		}
	}
	return inv
}

func dot8(a, b []float64) float64 {
	var s float64
	for i := 0; i < 8; i++ {
		s += a[i] * b[i]
	}
	return s
}