// Based on code by Matthias Müller (Ten Minute Physics), MIT License
// See licenses/TenMinutePhysics for full license text

package fluid

import (
	"math"
)

const (
	FluidDefaultDensity               = 1000 //kilogram per cubic meter
	FluidDefaultGravity               = -9.8
	FluidDefaultTimeStep              = 1. / 120
	FluidDefaultParticleRadiusPercent = 0.3
	FluidDefaultAngle                 = 0
	FluidIncompressibilityIterations  = 100
	FluidDefaultRestDensity           = 0
	FluidDefaultFlipPicRatio          = 0.9
	FluidDefaultOverRelaxation        = 1.9
	FluidDefaultSubSteps              = 2
)

type Fluid struct {
	Particles           []Particle
	Field               *VelocityField
	particleRadius      float32
	numCellParticles    []int     //SPATIAL HASHING: count of particles per cell
	firstCellParticle   []int     //SPATIAL HASHING: index of first particle for given cell
	cellParticleIds     []int     //SPATIAL HASHING: id of all particles
	pNumX, pNumY        int       //SPATIAL HASHING: number of particle cells x,y
	pNumCells           int       //SPATIAL HASHING: number of particle cells in total
	pInvSpacing         float32   //SPATIAL HASHING: one over the particle spacing
	particleDensity     []float32 //density of particles within each field cell
	particleRestDensity float32
	fluidDensity        float32
	angle               float32
	dt                  float32     //delta time step
	gravitynormal       FieldVector //unit vector for gravity
	stepacceleration    FieldVector //step acceleration due to gravity
	flipPicRatio        float32
	numSubSteps         int //number of simulation substeps
}

func NewFluid(dimensions FieldVector, spacing float32) *Fluid {
	f := &Fluid{
		Field:               NewVelocityField(dimensions, spacing),
		dt:                  FluidDefaultTimeStep,
		fluidDensity:        FluidDefaultDensity,
		particleRestDensity: FluidDefaultRestDensity,
		flipPicRatio:        FluidDefaultFlipPicRatio,
		numSubSteps:         FluidDefaultSubSteps,
	}

	f.Initialize()
	return f
}

func (f *Fluid) Initialize() {
	f.InitializeParticles()

	f.SetAngle(FluidDefaultAngle)
	f.particleDensity = make([]float32, f.Field.Nx*f.Field.Ny)
}

func (f *Fluid) InitializeParticles() {
	//compute particle radius based on our field spacing and radius scaling
	f.particleRadius = FluidDefaultParticleRadiusPercent * f.Field.H

	//we want to prepare a hexagonally packed grid of particles within our fluid field
	dx := float32(2.0 * f.particleRadius)
	dy := float32(math.Sqrt(3.) / 2. * float64(dx))

	//number of particles that fit within our specified dimensions
	numPx := int(math.Floor(float64(f.Field.Dimensions.X-2*f.Field.H-2*f.particleRadius) / float64(dx)))
	numPy := int(math.Floor(float64(f.Field.Dimensions.Y/2-2*f.Field.H-2*f.particleRadius) / float64(dy)))

	particlecount := numPx * numPy
	f.Particles = make([]Particle, particlecount)

	//now prepare our hashing structures, purely for collision adjustment
	//first let's compute some boundaries
	f.pInvSpacing = 1 / (2.2 * f.particleRadius)
	f.pNumX = int(math.Floor(float64(f.Field.Dimensions.X*f.pInvSpacing))) + 1
	f.pNumY = int(math.Floor(float64(f.Field.Dimensions.Y*f.pInvSpacing))) + 1
	f.pNumCells = f.pNumX * f.pNumY

	//now use boundaries to construct our spatial hashing containers
	f.numCellParticles = make([]int, f.pNumCells)
	f.firstCellParticle = make([]int, f.pNumCells+1)
	f.cellParticleIds = make([]int, int(particlecount))

	//create particles
	var p int = 0
	for i := 0; i < numPx; i++ {
		for j := 0; j < numPy; j++ {

			var rowAdjust float32 = 0
			if j%2 == 1 {
				rowAdjust = f.particleRadius
			}

			f.Particles[p].Position.X = f.Field.H + f.particleRadius + dx*float32(i) + rowAdjust
			f.Particles[p].Position.Y = f.Field.H + f.particleRadius + dy*float32(j)

			p++
		}
	}
}

func (f *Fluid) Step() {

	for range f.numSubSteps {
		f.IntegrateParticles()
		f.SeparateParticles()
		f.HandleParticleCollisions()
		f.TransferVelocities(true)
		f.UpdateParticleDensity()
		f.SolveIncompressibility()
		f.TransferVelocities(false)
	}
	//f.UpdateParticleColors()
}

func (f *Fluid) RecalculateStepAcceleration() {
	scale := FluidDefaultGravity * f.dt
	f.stepacceleration = FieldVector{
		X: scale * f.gravitynormal.X,
		Y: scale * f.gravitynormal.Y,
	}
}

// change the time step
func (f *Fluid) SetTimeStep(dt float32) {
	f.dt = dt
	f.RecalculateStepAcceleration()
}

// set new fluid angle, recompute gravity normal, initiate step accel recomputation
func (f *Fluid) SetAngle(angle float32) {
	f.angle = angle
	sin, cos := math.Sincos(float64(angle))
	f.gravitynormal = FieldVector{
		X: -float32(sin),
		Y: float32(cos),
	}
	f.RecalculateStepAcceleration()
}

func (f *Fluid) IntegrateParticles() {
	ax, ay := f.stepacceleration.X, f.stepacceleration.Y
	dt := f.dt
	for i := range f.Particles {
		p := &f.Particles[i]
		p.Velocity.X += ax
		p.Velocity.Y += ay
		p.Position.X += p.Velocity.X * dt
		p.Position.Y += p.Velocity.Y * dt
	}
}

// returns adjusted value between start and end, whichever is closer
// unless value is between, in which case it returns value
func clamp(value, start, end int) int {
	if value < start {
		return start
	}

	if value < end {
		return value
	}

	return end
}

func clamp32(value, start, end float32) float32 {
	if value < start {
		return start
	}

	if value < end {
		return value
	}

	return end
}

func (f *Fluid) PerformParticleSpatialHash() {

	clear(f.numCellParticles)

	//find spatcial hash cell for this particle's position, increment count there
	for i := range f.Particles {
		p := &f.Particles[i]
		xi := clamp(int(p.Position.X*f.pInvSpacing), 0, f.pNumX-1)
		yi := clamp(int(p.Position.Y*f.pInvSpacing), 0, f.pNumY-1)
		cellnumber := xi*f.pNumY + yi
		f.numCellParticles[cellnumber]++
	}

	//build up firstCellParticle bucket
	var sum int = 0
	for i := range f.pNumCells {
		sum += f.numCellParticles[i]
		f.firstCellParticle[i] = sum
	}
	f.firstCellParticle[f.pNumCells] = sum

	//next work backwards on it to assign particle ids, adjusting the
	//firstCellParticle bucket as we go
	for i := range f.Particles {
		p := &f.Particles[i]
		xi := clamp(int(p.Position.X*f.pInvSpacing), 0, f.pNumX-1)
		yi := clamp(int(p.Position.Y*f.pInvSpacing), 0, f.pNumY-1)
		cellnumber := xi*f.pNumY + yi
		f.firstCellParticle[cellnumber]--
		f.cellParticleIds[f.firstCellParticle[cellnumber]] = i
	}
}

func (f *Fluid) SeparateParticles() {

	//setup
	minDist := 2 * f.particleRadius
	minDist2 := minDist * minDist

	//build our spatial hashing table
	f.PerformParticleSpatialHash()

	//perform broad phase search
	for i := range f.Particles {

		p := &f.Particles[i]

		pxi := int(p.Position.X * f.pInvSpacing)
		pyi := int(p.Position.Y * f.pInvSpacing)
		x0 := max(pxi-1, 0)
		y0 := max(pyi-1, 0)
		x1 := min(pxi+1, f.pNumX)
		y1 := min(pyi+1, f.pNumY)

		for xi := x0; xi < x1; xi++ {
			for yi := y0; yi < y1; yi++ {
				//find current cell
				cell := xi*f.pNumY + yi

				//set start and ending particles to check between this and next cell
				first := f.firstCellParticle[cell]
				last := f.firstCellParticle[cell+1]

				//iterate over these candidates
				for j := first; j < last; j++ {
					//skip self
					id := f.cellParticleIds[j]
					if id == i {
						continue
					}

					//compute delta
					dx := f.Particles[id].Position.X - p.Position.X
					dy := f.Particles[id].Position.Y - p.Position.Y
					dist2 := dx*dx + dy*dy

					//rule out non-collisions, ignore overlaps (we screwed)
					if (dist2 == 0) || (dist2 > minDist2) {
						continue
					}

					//need to use more expensive sqrt now to find our penetration
					//then compute a scalar for vector normalization, applied
					//equally in half to both particles
					dist := float32(math.Sqrt(float64(dist2)))
					penetration := (minDist - dist)
					scale := penetration / dist
					dx *= scale / 2
					dy *= scale / 2

					p.Position.X -= dx
					p.Position.Y -= dy
					f.Particles[id].Position.X += dx
					f.Particles[id].Position.Y += dy
				}
			}
		}
	}
}

func (f *Fluid) HandleParticleCollisions() {
	//setup
	/*
		minDist := 2 * f.particleRadius
		minDist2 := minDist * minDist
	*/

	minX := f.particleRadius
	maxX := float32(f.Field.Nx-1)*f.Field.H - f.particleRadius
	minY := f.particleRadius
	maxY := float32(f.Field.Ny-1)*f.Field.H - f.particleRadius

	for i := range f.Particles {
		p := &f.Particles[i]

		if p.Position.X < minX {
			p.Position.X = minX
			p.Velocity.X = 0
		}

		if p.Position.X > maxX {
			p.Position.X = maxX
			p.Velocity.X = 0
		}

		if p.Position.Y < minY {
			p.Position.Y = minY
			p.Velocity.Y = 0
		}

		if p.Position.Y > maxY {
			p.Position.Y = maxY
			p.Velocity.Y = 0
		}

	}
}

func (f *Fluid) TransferVelocities(togrid bool) {

	//setup
	n := f.Field.Ny
	h1 := f.Field.InvH
	h2 := 0.5 * f.Field.H

	if togrid {
		copy(f.Field.PrevU, f.Field.U)
		copy(f.Field.PrevV, f.Field.V)
		clear(f.Field.DU)
		clear(f.Field.DV)
		clear(f.Field.U)
		clear(f.Field.V)

		for i := 0; i < f.Field.Numcells; i++ {
			if f.Field.S[i] == 0 {
				f.Field.CellType[i] = CellTypeSolid
			} else {
				f.Field.CellType[i] = CellTypeAir
			}
		}

		for i := range f.Particles {
			x := f.Particles[i].Position.X
			y := f.Particles[i].Position.Y
			xi := clamp(int(math.Floor(float64(x*h1))), 0, f.Field.Nx-1)
			yi := clamp(int(math.Floor(float64(y*h1))), 0, f.Field.Ny-1)
			cellnumber := xi*n + yi
			if f.Field.CellType[cellnumber] == CellTypeAir {
				f.Field.CellType[cellnumber] = CellTypeFluid
			}
		}
	}

	for component := 0; component < 2; component++ {
		var dx, dy float32
		var fl, prevF, d *[]float32

		if component == 0 {
			dx = 0
			dy = h2
			fl = &f.Field.U
			prevF = &f.Field.PrevU
			d = &f.Field.DU
		} else {
			dx = h2
			dy = 0
			fl = &f.Field.V
			prevF = &f.Field.PrevV
			d = &f.Field.DV
		}

		for i := range f.Particles {
			p := &f.Particles[i]
			x := p.Position.X
			y := p.Position.Y

			x = clamp32(x, f.Field.H, float32((f.Field.Nx-1))*f.Field.H)
			y = clamp32(y, f.Field.H, float32((f.Field.Ny-1))*f.Field.H)

			//find cell components neighbouring x,y
			x0 := min(int(math.Floor(float64((x-dx)*h1))), f.Field.Nx-2)
			x1 := min(x0+1, f.Field.Nx-2)
			y0 := min(int(math.Floor(float64((y-dy)*h1))), f.Field.Ny-2)
			y1 := min(y0+1, f.Field.Ny-2)

			//compute scales
			tx := ((x - dx) - float32(x0)*f.Field.H) * h1
			ty := ((y - dy) - float32(y0)*f.Field.H) * h1
			sx := 1.0 - tx
			sy := 1.0 - ty

			//compute weights
			d0 := sx * sy
			d1 := tx * sy
			d2 := tx * ty
			d3 := sx * ty

			//cell indexes
			nr0 := x0*n + y0
			nr1 := x1*n + y0
			nr2 := x1*n + y1
			nr3 := x0*n + y1

			if togrid {
				//extract velocity for this component
				var pv float32 // = this.particleVel[2 * i + component];
				if component == 0 {
					pv = p.Velocity.X
				} else {
					pv = p.Velocity.Y
				}

				//to our velocity component we will add the contribution of
				//the particle velocity scaled according to its weight
				(*fl)[nr0] += pv * d0
				(*fl)[nr1] += pv * d1
				(*fl)[nr2] += pv * d2
				(*fl)[nr3] += pv * d3

				//we also keep a running total of all the weights (d is r in the notes)
				(*d)[nr0] += d0
				(*d)[nr1] += d1
				(*d)[nr2] += d2
				(*d)[nr3] += d3
			} else {

				var offset int
				var v float32
				if component == 0 {
					offset = n
					v = p.Velocity.X
				} else {
					offset = 1
					v = p.Velocity.Y
				}

				var valid0, valid1, valid2, valid3 float32 = 0., 0., 0., 0.

				if f.Field.CellType[nr0] != CellTypeAir || f.Field.CellType[nr0-offset] != CellTypeAir {
					valid0 = 1
				}

				if f.Field.CellType[nr1] != CellTypeAir || f.Field.CellType[nr1-offset] != CellTypeAir {
					valid1 = 1
				}

				if f.Field.CellType[nr2] != CellTypeAir || f.Field.CellType[nr2-offset] != CellTypeAir {
					valid2 = 1
				}

				if f.Field.CellType[nr3] != CellTypeAir || f.Field.CellType[nr3-offset] != CellTypeAir {
					valid3 = 1
				}

				//weights
				wsum := valid0*d0 + valid1*d1 + valid2*d2 + valid3*d3

				if wsum > 0.0 {
					picV := (valid0*d0*(*fl)[nr0] + valid1*d1*(*fl)[nr1] + valid2*d2*(*fl)[nr2] + valid3*d3*(*fl)[nr3]) / wsum
					corr := (valid0*d0*((*fl)[nr0]-(*prevF)[nr0]) + valid1*d1*((*fl)[nr1]-(*prevF)[nr1]) +
						valid2*d2*((*fl)[nr2]-(*prevF)[nr2]) + valid3*d3*((*fl)[nr3]-(*prevF)[nr3])) / wsum
					flipV := v + corr
					vNew := (1-f.flipPicRatio)*picV + f.flipPicRatio*flipV

					if component == 0 {
						p.Velocity.X = vNew
					} else {
						p.Velocity.Y = vNew
					}
				}
			}
		}

		if togrid {
			//finally, for all q's, we divide by the weighted sum of that cell
			for i := 0; i < len(*fl); i++ {
				if (*d)[i] > 0 {
					(*fl)[i] = (*fl)[i] / (*d)[i]
				}
			}

			//next we restore cells are solid to their previous values
			for i := 0; i < f.Field.Nx; i++ {
				for j := 0; j < f.Field.Ny; j++ {
					idx := i*n + j
					issolid := f.Field.CellType[idx] == CellTypeSolid

					if issolid || (i > 0 && f.Field.CellType[(i-1)*n+j] == CellTypeSolid) {
						f.Field.U[idx] = f.Field.PrevU[idx]
					}

					if issolid || (j > 0 && f.Field.CellType[i*n+j-1] == CellTypeSolid) {
						f.Field.V[idx] = f.Field.PrevV[idx]
					}

				}
			}
		}

	}

}

func (f *Fluid) UpdateParticleDensity() {
	n := f.Field.Ny
	h := f.Field.H
	h1 := f.Field.InvH
	h2 := 0.5 * h

	clear(f.particleDensity)

	for i := range f.Particles {
		p := &f.Particles[i]
		x := p.Position.X
		y := p.Position.Y

		x = clamp32(x, h, float32(f.Field.Nx-1)*h)
		y = clamp32(y, h, float32(f.Field.Ny-1)*h)

		x0 := int(math.Floor(float64((x - h2) * h1)))
		x1 := min(x0+1, f.Field.Nx-2)

		y0 := int(math.Floor(float64((y - h2) * h1)))
		y1 := min(y0+1, f.Field.Ny-2)

		tx := ((x - h2) - float32(x0)*h) * h1
		ty := ((y - h2) - float32(y0)*h) * h1

		sx := 1.0 - tx
		sy := 1.0 - ty

		if x0 < f.Field.Nx && y0 < f.Field.Ny {
			f.particleDensity[x0*n+y0] += sx * sy
		}
		if x1 < f.Field.Nx && y0 < f.Field.Ny {
			f.particleDensity[x1*n+y0] += tx * sy
		}
		if x1 < f.Field.Nx && y1 < f.Field.Ny {
			f.particleDensity[x1*n+y1] += tx * ty
		}
		if x0 < f.Field.Nx && y1 < f.Field.Ny {
			f.particleDensity[x0*n+y1] += sx * ty
		}
	}

	if f.particleRestDensity == 0.0 {
		var sum float32 = 0.0
		numFluidCells := 0

		for i := 0; i < f.Field.Nx*f.Field.Ny; i++ {
			if f.Field.CellType[i] == CellTypeFluid {
				sum += f.particleDensity[i]
				numFluidCells++
			}
		}

		if numFluidCells > 0 {
			f.particleRestDensity = sum / float32(numFluidCells)
		}
	}
}

func (f *Fluid) SolveIncompressibility() {

	//setup
	copy(f.Field.PrevU, f.Field.U)
	copy(f.Field.PrevV, f.Field.V)

	n := f.Field.Ny
	//cp := f.fluidDensity * f.Field.H / FluidDefaultTimeStep

	for iter := 0; iter < FluidIncompressibilityIterations; iter++ {
		for i := 1; i < f.Field.Nx-1; i++ {
			for j := 1; j < f.Field.Ny-1; j++ {

				if f.Field.CellType[i*n+j] != CellTypeFluid {
					continue
				}

				//compute indexes of surrounding cells
				center := i*n + j
				left := (i-1)*n + j
				right := (i+1)*n + j
				bottom := i*n + j - 1
				top := i*n + j + 1

				sx0 := f.Field.S[left]
				sx1 := f.Field.S[right]
				sy0 := f.Field.S[bottom]
				sy1 := f.Field.S[top]
				s := sx0 + sx1 + sy0 + sy1
				if s == 0 {
					continue
				}

				//divergence -> we want this to be zero
				div := f.Field.U[right] - f.Field.U[center] + f.Field.V[top] - f.Field.V[center]

				if f.particleRestDensity > 0.0 {
					var k float32 = 1.0
					var compression float32 = f.particleDensity[i*n+j] - f.particleRestDensity
					if compression > 0.0 {
						div = div - k*compression
					}
				}

				p := -div / s
				p *= FluidDefaultOverRelaxation //overRelaxation
				//f.Field.Pressure[center] += cp * p

				f.Field.U[center] -= sx0 * p
				f.Field.U[right] += sx1 * p
				f.Field.V[center] -= sy0 * p
				f.Field.V[top] += sy1 * p

			}
		}
	}

}