// A particle creates an attractive force field towards it s.t. the strength is proportional to the particle's mass.
// It doesn't depend on the distance - the force is constant.

import { Arithmetic, Vector, Scalar } from "../arithmetic";
import { Duration, MicroDuration, random, System } from "../helpers";

export type Config = {
  dt: MicroDuration, // e.g. 5 miliseconds
  // Box in which to generate random particles with uniform distribution, and random velocities
  topLeft: {x: number, y: number},
  bottomRight: {x: number, y: number},
  maxSpeed: number,

  countSmall: number,
  massRangeSmall: { min: number, max: number },
  countBig: number,
  massRangeBig: { min: number, max: number },
}

export type State = {
  position: Vector,
  velocity: Vector,
  mass: Scalar,
}

export type Observables = {
  particleCount: number,
  averageSpeed: number,
  netMass: number,
  netForce: { fx: number, fy: number },
  netMomentum: { px: number, py: number },
  centerOfMass: { x: number, y: number },
  netKineticEnergy: number,
  netPotentialEnergy: number,
  netPotentialEnergyIntegralCalculation: number,
  netEnergy: number,
}

export type Capabilities = {
  spawn(x: number, y: number, vx: number, vy: number, m: number): void,
}

export type FlatGravitySystem = System<State, Observables, Capabilities>;

export function FlatGravitySystem(config: Config): FlatGravitySystem {
  const arith = Arithmetic({ count: 0, capacity: 10000 });
  const { Vector, Scalar } = arith;

  // Responsible for computing the force/acceleration fields
  function Force() {
    const force = Vector.allocate();
    const acceleration = Vector.allocate();

    function LOOM() {
      const count = arith.getCount();

      Vector.zero(force);
      Vector.zero(acceleration);
      for (let i = 0; i < count; i++) {
        // Compute the net acceleration on i coming from interaction with all the other particles.
        for (let j = i + 1; j < count; j++) {
          const mi = mass.val[i];
          const mj = mass.val[j];
          const dx = position.x[j] - position.x[i];
          const dy = position.y[j] - position.y[i];
          const size = Math.sqrt(dx*dx + dy*dy);
          if (size == 0) { continue } 
          const ux = dx/size;
          const uy = dy/size;
          // [ux, uy] is the unit vector pointing from i to j

          const ax_i = mj*G*ux;
          const ay_i = mj*G*uy;
          const fx_i = mi*ax_i;
          const fy_i = mi*ay_i;
          const ax_j = -mi*G*ux;
          const ay_j = -mi*G*uy;

          acceleration.x[i] += ax_i;
          acceleration.y[i] += ay_i;
          acceleration.x[j] += ax_j;
          acceleration.y[j] += ay_j;

          force.x[i] +=  fx_i;
          force.y[i] +=  fy_i;
          force.x[j] += -fx_i;
          force.y[j] += -fy_i;
        }
      }
    }

    return { force, acceleration, LOOM };
  }

  // TODO: Better name
  // Responsible for stepping (position, velocity)
  function MicroEvolution() {
    const microDisplacement: Vector = Vector.allocate();

    function LOOM(dt: MicroDuration) {
      // Symplectic Euler:
      // velocity = velocity + dt*acceleration
      // position = position + dt*velocity
      //
      // Explicit Euler (seems to increase total energy over time):
      // position = position + dt*velocity
      // velocity = velocity + dt*acceleration

      Vector.addScaledByConstant(velocity, velocity, dt, acceleration);

      // If you don't need micro-displacement, you can just use:
      // Vector.addScaledByConstant(position, position, dtInSeconds, velocity);
      Vector.scaleByConstant(microDisplacement, velocity, dt);
      Vector.add(position, position, microDisplacement);
    }
    return { microDisplacement, LOOM };
  }
  
  // Updates potential energy by current forces
  function PotentialEnergyUpdate() {
    const potentialEnergy: Scalar = Scalar.allocate();
    const microWork: Scalar = Scalar.allocate();

    function LOOM(force: Vector, microDisplacement: Vector) {
      // U := U - <F, microDisplacement>
      Vector.innerProduct(microWork, force, microDisplacement);
      Scalar.addScaledByConstant(potentialEnergy, potentialEnergy, -1, microWork);
    }
    return { potentialEnergy, LOOM };
  }

  function NetKineticEnergy() {
    const kineticEnergy: Scalar = Scalar.allocate();
    const velocitySquared: Scalar = Scalar.allocate();
    const netKineticEnergy = { value: 0 };

    function LOOM(mass: Scalar, velocity: Vector) {
      // totalKineticEnergy = 1/2 * sum(m*v^2)
      Vector.innerProduct(velocitySquared, velocity, velocity);
      Scalar.mul(kineticEnergy, mass, velocitySquared);
      netKineticEnergy.value = 0.5*Scalar.sum(kineticEnergy);
    }
    return { netKineticEnergy, LOOM };
  }

  function SetMomentumAndCenterOfMass() {
    const momentum: Vector = Vector.allocate();
    const weightedPosition: Vector = Vector.allocate();
    const centerOfMass = { x: 0, y: 0 };
    const netMass = { value: 0 };

    function LOOM(mass: Scalar, velocity: Vector) {
      Vector.scaleByScalar(momentum, velocity, mass);

      netMass.value = Scalar.sum(mass);

      // centerOfMass := sum(mass * position) / totalMass
      Vector.scaleByScalar(weightedPosition, position, mass);
      const { x: mx, y: my } = Vector.sum(weightedPosition);
      if (netMass.value !== 0) {
        centerOfMass.x = mx/netMass.value;
        centerOfMass.y = my/netMass.value;
      }
    }
    return { momentum, centerOfMass, LOOM }
  }

  let time: Duration = 0;

  const state: State = {
    position: Vector.allocate(), // in pixels
    velocity: Vector.allocate(), // in pixels per second (not miliseconds)
    mass: Scalar.allocate(),
  };
  const { velocity, position, mass } = state;

  const G = 0.05;

  const { force, acceleration, LOOM: forceLoom } = Force();
  const { microDisplacement, LOOM: microEvolutionLoom } = MicroEvolution();
  const { potentialEnergy, LOOM: potentialEnergyUpdateLoom } = PotentialEnergyUpdate();
  const { netKineticEnergy, LOOM: netKineticEnergyLoom }  = NetKineticEnergy();
  const { momentum, centerOfMass, LOOM: setMomentumAndCenterOfMassLoom } = SetMomentumAndCenterOfMass();
  const weightedPosition: Vector = Vector.allocate();
  const speed: Scalar = Scalar.allocate();

  function spawn(x: number, y: number, vx: number, vy: number, m: number) {
    const idx = arith.getCount();

    position.x[idx] = x;
    position.y[idx] = y;
    velocity.x[idx] = vx;
    velocity.y[idx] = vy;
    state.mass.val[idx] = m;

    arith.setCount(idx + 1);
  }

  function reset() {
    arith.setCount(0);

    for (let i = 0; i < config.countBig; i++) {
      spawn(
        // Random position inside the bounding box
        random(config.topLeft.x, config.bottomRight.x),
        random(config.topLeft.y, config.bottomRight.y),
        // Random velocity between -maxVelocity and +maxVelocity
        random(-1, 1) * config.maxSpeed,
        random(-1, 1) * config.maxSpeed,
        // Random mass
        random(config.massRangeBig.min, config.massRangeBig.max),
      );
    }
    for (let i = 0; i < config.countSmall; i++) {
      spawn(
        // Random position inside the bounding box
        random(config.topLeft.x, config.bottomRight.x),
        random(config.topLeft.y, config.bottomRight.y),
        // Random velocity between -maxVelocity and +maxVelocity
        random(-1, 1) * config.maxSpeed,
        random(-1, 1) * config.maxSpeed,
        // Random mass
        random(config.massRangeSmall.min, config.massRangeSmall.max),
      );
    }

    Scalar.zero(potentialEnergy);
  }
  reset();

  // TODO: how to initialize this without code-duplication?
  const observables: Observables = {
    particleCount: 0,
    averageSpeed: 0,
    netMass: 0,
    netForce: { fx: 0, fy: 0 },
    netMomentum: { px: 0, py: 0 },
    centerOfMass: { x: 0, y: 0 },
    netKineticEnergy: 0,
    netPotentialEnergy: 0,
    netPotentialEnergyIntegralCalculation: 0,
    netEnergy: 0,
  };


  function refreshObservables() {
    const count = arith.getCount();
    observables.particleCount = count;

    if (count === 0) return;

    Vector.magnitude(speed, velocity);
    // Note that this would produce division by zero for `count === 0`.
    observables.averageSpeed = Scalar.average(speed);

    // netMass := sum(mass)
    const netMass = Scalar.sum(mass);
    observables.netMass = netMass;

    // netForce := sum(force)
    const { x: fx, y: fy } = Vector.sum(force);
    observables.netForce.fx = fx;
    observables.netForce.fy = fy;
    
    setMomentumAndCenterOfMassLoom(mass, velocity);
    const { x: px, y: py } = Vector.sum(momentum);
    observables.netMomentum.px = px;
    observables.netMomentum.py = py;

    observables.centerOfMass.x = centerOfMass.x;
    observables.centerOfMass.y = centerOfMass.y;

    netKineticEnergyLoom(mass, velocity);
    observables.netKineticEnergy = netKineticEnergy.value;

    // Potential energy
    let totalPotential = 0;
    for (let i = 0; i < count; i++) {
      for (let j = i + 1; j < count; j++) {
        const dx = position.x[j] - position.x[i];
        const dy = position.y[j] - position.y[i];
        const distance = Math.sqrt(dx * dx + dy * dy);
        
        // U = G * m1 * m2 * distance
        totalPotential += G * mass.val[i] * mass.val[j] * distance;
      }
    }
    observables.netPotentialEnergy = totalPotential;
    observables.netEnergy = observables.netKineticEnergy + observables.netPotentialEnergy;
    observables.netPotentialEnergyIntegralCalculation = Scalar.sum(potentialEnergy);
  }

  function microFlow(dt: MicroDuration) {
    forceLoom();
      const dtInSeconds = dt / 1000.0;
    microEvolutionLoom(dtInSeconds);
    potentialEnergyUpdateLoom(force, microDisplacement);
  }

  // Suppose that we need to flow for 43 miliseconds, and dt is 10 miliseconds.
  // Then we do micro-flow 4 times totalling 40 miliseconds. We don't do micro-flow for 3 seconds
  // - that could cause problems with integration.
  // Instead we remember this leftover time, so the next time we flow, we add this leftover time to the total time to flow.
  let leftover_time: Duration = 0;
  // Integrates the micro-flow, updates global time, and recomputes the observables.
  function flow(t: Duration) {
    // TODO: They way we compute leftover time is not so great. There must be a better way.
    const end_time = time + t + leftover_time;

    while (true) {
      const next_time = time + config.dt;
      if (next_time <= end_time) {
        time = next_time;
        microFlow(config.dt);
      } else {
        // assume at the start
        //   dt = 10 ms
        //
        //   time = 0
        //   leftover_time = 0
        //   t = 43 ms
        // then here
        //   next_time = 50 ms
        //   time      = 40 ms
        //   end_time  = 43 ms
        // so actually we have 3 miliseconds leftover, i.e. we set
        leftover_time = end_time - time;
        break;
      }
    }

  }

  return {
    state,
    observables,
    flow,
    refreshObservables,
    totalTime() { return time; },
    reset,
    capabilities: {
      spawn,
    },
  };
}

