/**
 *
 * @source: https://software.mcgoron.com/peter/scattering.js
 *
 * @licstart  The following is the entire license notice for the 
 *  JavaScript code in this file.
 *
 * Copyright (C) 2023 Peter McGoron
 *
 * The JavaScript code in this file is free software: you can
 * redistribute it and/or modify it under the terms of the GNU
 * General Public License (GNU GPL) as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option)
 * any later version.  The code is distributed WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE.  See the GNU GPL for more details.
 *
 * As additional permission under GNU GPL version 3 section 7, you
 * may distribute non-source (e.g., minimized or compacted) forms of
 * that code without the copy of the GNU GPL normally required by
 * section 4, provided you include this license notice and a URL
 * through which recipients can access the Corresponding Source.
 *
 * @licend  The above is the entire license notice
 * for the JavaScript code in this page.
 *
 */

/*****************************************
 * Math
 ****************************************/

function sign(x) {
  if (x < 0) return -1;
  if (x >= 0) return 1;
}

function distance(x,y) {
  return Math.sqrt(x**2 + y**2);
}

function veryClose(x, y) {
  return Math.abs(x - y) < 0.001;
}

/*****************************************
 * Object Classes
 ****************************************/

class Rect {
  constructor(args) {
    if ('r1' in args) {
      this.x_0 = Math.min(args.r1.x_0, args.r2.x_0);
      this.y_0 = Math.min(args.r1.y_0, args.r2.y_0);
      this.x_1 = Math.max(args.r1.x_1, args.r2.x_1);
      this.y_1 = Math.max(args.r1.y_1, args.r2.y_1);
    } else {
      this.x_0 = args.x_0;
      this.x_1 = args.x_1;
      this.y_0 = args.y_0;
      this.y_1 = args.y_1;
    }
  }

  containedIn(rect) {
    return this.x_0 >= rect.x_0 && this.x_1 <= rect.x_1
        && this.y_0 >= rect.y_0 && this.y_1 <= rect.y_1;
  }

  area() {
    return (this.x_1 - this.x_0)*(this.y_1 - this.y_0);
  }
}

let windowRect = new Rect({x_0: 0, x_1: 800, y_0: 0, y_1: 600})
let width = windowRect.x_1;
let height = windowRect.y_1;

/* Infinite lines are drawn with length = 2*longestDiag which will always
 * cross the entire screen. */
let longestDiag = distance(width, height);

/* Basic infinite line class, useful for physics simulation */
class InfiniteLine {
  constructor(x_0, y_0, angle) {
    this.x_0 = x_0;
    this.y_0 = y_0;
    this.angle = angle;
  }

  insideOfRect(r) {
    // implicit equation of a line:
    // (y - y_0)cos(theta) - (x - x_0)*sin(theta) = 0
    // Evaluate the LHS for two points of a line segment. If the
    // signs differ, then the two segments intersect.
    let f = (x,y) => (
      (y - this.y_0)*Math.cos(this.angle) - (x - this.x_0)*Math.sin(this.angle)
    )

    const upper_left = f(0,0);
    const upper_right = f(width,0);
    if (sign(upper_left) != sign(upper_right))
      return true;

    const lower_left = f(0,height);
    if (sign(lower_left) != sign(upper_left))
      return true;

    const lower_right = f(width,height);
    return sign(lower_right) != sign(lower_left)
        || sign(lower_right) != sign(upper_right);
  }

  /* If the line is outside of the canvas */
  cullable() {
    return !this.insideOfRect(windowRect);
  }

  // Returns true if drawn.
  draw() {
    if (this.cullable())
      return false;

    fill(color(0,0,0));
    stroke(color(0,0,0));
    translate(this.x_0, this.y_0);
    rotate(this.angle);
    strokeWeight(2);
    line(-longestDiag, 0, longestDiag, 0);
    resetMatrix();
    return true;
  }
}

class PlaneWave extends InfiniteLine {
  constructor(x_0, y_0, angle, v) {
    super(x_0, y_0, angle);
    this.v = v;
    this.atomsHit = new Set();
  }

  update() {
    this.x_0 = this.x_0 + -this.v * deltaTime * Math.sin(this.angle);
    this.y_0 = this.y_0 + this.v * deltaTime * Math.cos(this.angle);
  }

  drawArrow() {
    fill(color(0,0,0));
    stroke(color(0,0,0));
    translate(this.x_0, this.y_0);
    rotate(this.angle);
    triangle(-5, 0, 0, 15, 5, 0);
    resetMatrix();
  }
}

class SphericalWave {
  constructor (x,y) {
    this.x = x;
    this.y = y;
    this.ul_dist = distance(this.x, this.y);
    this.ur_dist = distance(width - this.x, this.y);
    this.ll_dist = distance(this.x, this.y - height);
    this.lr_dist = distance(this.x - width, this.y - height);
    this.rad = 1;
  }

  cullable() {
    return this.ul_dist < this.rad && this.ur_dist < this.rad
        && this.ll_dist < this.rad && this.lr_dist < this.rad;
  }

  update() {
    this.rad += 0.05*deltaTime;
  }

  draw() {
    stroke(color(255,255,0));
    strokeWeight(2);
    noFill();
    circle(this.x, this.y, 2*this.rad);
  }
}

class Atom {
  static radius = 5;
  constructor(x,y) {
    this.x = x;
    this.y = y;
    this.color = color(255,0,0);
    this.wasDrawn = false;
    this.boundRect = new Rect({
      x_0: x - Atom.radius,
      x_1: x + Atom.radius,
      y_0: y - Atom.radius,
      y_1: y + Atom.radius
    });
  }

  collidesWithLine(line) {
    if (line.atomsHit.has(this))
      return false;
    // Calculate vector from point on line to center of atom
    const atom_vec_x = - line.x_0 + this.x;
    const atom_vec_y = - line.y_0 + this.y;
    // Unit vector of the line is (cos(theta),sin(theta))
    // calculate dot product of 
    const dot_prod = Math.cos(line.angle)*atom_vec_x
                     + Math.sin(line.angle)*atom_vec_y;

    const proj_x = line.x_0 + dot_prod*Math.cos(line.angle);
    const proj_y = line.y_0 + dot_prod*Math.sin(line.angle);

    return distance(proj_x - this.x, proj_y - this.y) <= Atom.radius;
  }

  updatePos(x,y) {
    this.x = x;
    this.y = y;
    this.boundRect.x_0 = x - Atom.radius;
    this.boundRect.x_1 = x + Atom.radius;
    this.boundRect.y_0 = y - Atom.radius;
    this.boundRect.y_1 = y + Atom.radius;
  }

  draw(_color) {
    this.wasDrawn = true;
    if (_color === undefined)
      fill(this.color);
    else
      fill(_color);
    stroke(color(0,0,0));
    strokeWeight(1);
    circle(this.x, this.y, 2*Atom.radius);
  }
}

class AABBNode {
  constructor(args) {
    this.leftNode = null;
    this.rightNode = null;
    if (args !== undefined) {
      this.obj = args.obj;
      this.rect = args.rect;
    } else {
      this.obj = null;
      this.rect = null;
    }
  }

  _lineCollisions(line, arr) {
    if (this.obj !== null) {
      if (this.obj.collidesWithLine(line)) {
        arr.push(this.obj);
        return null;
      }
    }

    if (this.leftNode && line.insideOfRect(this.leftNode.rect))
      this.leftNode._lineCollisions(line, arr);
    if (this.rightNode && line.insideOfRect(this.rightNode.rect))
      this.rightNode._lineCollisions(line, arr);
  }

  lineCollisions(line) {
    let arr = [];
    this._lineCollisions(line, arr);
    return arr;
  }

  add(obj, rect) {
    // Empty node (the root). Add object.
    if (this.obj === null && this.leftNode === null && this.rightNode === null) {
      this.obj = obj;
      this.rect = rect;
    } else if (this.obj !== null) {
      // Leaf node. Split node.
      this.leftNode = new AABBNode({obj: this.obj, rect: this.rect});
      this.rightNode = new AABBNode({obj: obj, rect: rect});
      this.rect = new Rect({r1: this.rect, r2: rect});
      this.obj = null;
    } else {
      // General case with two children.
      // If the rectangle is entirely contained on one node, put the
      // rectangle into that node.

      if (rect.containedIn(this.leftNode.rect)) {
        this.leftNode.add(obj, rect);
      } else if (rect.containedIn(this.rightNode.rect)) {
        this.rightNode.add(obj, rect);
      } else {

        // Otherwise, calculate the node that will add the least amount of
        // area to the tree.
        const newLeftRect = new Rect({r1: this.leftNode.rect, r2: rect});
        const newRightRect = new Rect({r1: this.rightNode.rect, r2: rect});

        if (newLeftRect.area() < newRightRect.area()) {
          this.leftNode.add(obj, rect);
        } else {
          this.rightNode.add(obj, rect);
        }

        this.rect = new Rect({r1 : this.leftNode.rect, r2 : this.rightNode.rect})
      }
    }
  }

  draw(n) {
    if (this.rect !== null) {
      fill(color(0, n, 0, 100));
      quad(this.rect.x_0, this.rect.y_0,
           this.rect.x_0, this.rect.y_1,
           this.rect.x_1, this.rect.y_1,
           this.rect.x_1, this.rect.y_0);
    }

    if (this.obj !== null) {
      this.obj.draw();
    }

    if (this.leftNode) {
      this.leftNode.draw(n+10);
    }

    if (this.rightNode) {
      this.rightNode.draw(n+10);
    }
  }
}

/********************************
 * Mouse handler classes.
 *
 * The mouse handler functions return either ``null`` or an object
 * created by ``stateTo``, which contains the name of the new state
 * and the return value of the handler function.
 *******************************/

// Base class with dummy functions.
class MouseHandler {
  stateTo(newState, rval) {
    return {newState: newState, rval: rval};
  }
  handleMouseHold() {return null;}
  handleMouseUp() {return null;}
  handleMouseDown() {return null;}
  handleMouseMove() {return null;}
}

class PlaneWaveMouseHandler extends MouseHandler {
  constructor() {
    super();
    this.curWave = null;
    this.oldAngle = -Math.PI/4;
  }

  handleMouseDown() {
    this.curWave = new PlaneWave(mouseX, mouseY, this.oldAngle, 0.1);
    return null;
  }

  handleMouseMove() {
    if (!this.curWave) return null;
    /* Change angle to where the mouse cursor is pointing */
    if (!veryClose(mouseX, this.curWave.x_0)  ||
        !veryClose(mouseY, this.curWave.y_0)) {
        this.oldAngle = Math.atan2( -(mouseX - this.curWave.x_0), mouseY - this.curWave.y_0);
        this.curWave.angle = this.oldAngle;
    }

    /* TODO: make faster plane waves? */

    this.curWave.draw();
    this.curWave.drawArrow();
    return null;
  }

  handleMouseUp() {
    const oldWave = this.curWave;
    this.curWave = null;
    return this.stateTo('base', oldWave);
  }
}

class AtomMouseHandler extends MouseHandler {
  constructor() {
    super();
    this.currentAtom = null;
  }

  handleMouseDown() {
    if (!this.currentAtom) {
      this.currentAtom = new Atom(mouseX, mouseY);
      this.currentAtom.draw();
      return null;
    } else {
      const oldAtom = this.currentAtom;
      this.currentAtom = null;
      return this.stateTo('base', oldAtom);
    }
  }

  handleMouseUp() {
    return this.handleMouseMove();
  }

  handleMouseHold() {
    return this.handleMouseMove();
  }

  handleMouseMove() {
    this.currentAtom.updatePos(mouseX, mouseY);
    this.currentAtom.draw();
    return null;
  }
}

/* This class is the default mouse handler. */
class SimulationMouseHandler extends MouseHandler {
  constructor() {
    super();
    this.putWave = new PlaneWaveMouseHandler();
    this.putAtom = new AtomMouseHandler();
    this.base = this;

    this.mouseDown = false;
    this.state = "base";
  }

  handleMouseDown() {
      console.log(keyCode);
      if (keyIsPressed && keyCode == CONTROL) {
        this.state = 'putAtom';
      } else {
        this.state = 'putWave';
      }

      return this[this.state].handleMouseDown();
  }

  /* Function called by the simulation to handle dispatch */
  handleMouse() {
    var r;

    if (mouseX < 0 || mouseX > width || mouseY < 0 || mouseY > height)
      return null;

    if (mouseIsPressed && !this.mouseDown) {
      this.mouseDown = true;
      r = this[this.state].handleMouseDown();
    } else if (!mouseIsPressed) {
      this.mouseDown = false;
      r = this[this.state].handleMouseUp();
    } else {
      r = this[this.state].handleMouseMove();
    }

    if (r == null)
      return null;
    this.state = r.newState;

    return r.rval;
  }
}

/***************************
 * Simulation handler
 **************************/

class Simulation {
  constructor() {
    this.planeWaves = [];
    this.sphericalWaves = [];
    this.atomTree = new AABBNode();
    this.atoms = [];
    this.paused = false;
    this.enterKeyPressed = false;
    this.mouseHandler = new SimulationMouseHandler();

    const _this = this;
    document.getElementById("clear_waves").addEventListener("click", function (e) {
      _this.planeWaves = [];
      _this.sphericalWaves = [];
    });
    document.getElementById("clear_atoms").addEventListener("click", function (e) {
      _this.atoms = [];
      _this.atomTree = new AABBNode();
    });
  }

  addAtom(atom) {
    this.atomTree.add(atom, atom.boundRect);
    this.atoms.push(atom);
  }

  bragg_setup() {
    for (let x = 50; x < width; x += 50) {
      for (let y = height - 200; y < height; y += 50) {
        this.addAtom(new Atom(x, y));
      }
    }
  }

  update() {
    let handleRet = this.mouseHandler.handleMouse();
    if (handleRet instanceof PlaneWave) {
      this.planeWaves.push(handleRet);
    } else if (handleRet instanceof Atom) {
      this.addAtom(handleRet);
    }

    if (!this.enterKeyPressed && keyIsPressed && keyCode == ENTER) {
      this.enterKeyPressed = true;
      this.paused = !this.paused;
    } else {
      this.enterKeyPressed = false;
    }

    for (let ind = 0; ind < this.planeWaves.length; ind++) {
      if (!this.paused)
        this.planeWaves[ind].update();
      if (!this.planeWaves[ind].draw()) {
        this.planeWaves.splice(ind--, 1);
        continue;
      }

      if (!this.paused) {
        let newarr = this.atomTree.lineCollisions(this.planeWaves[ind]);
        for (let j = 0; j < newarr.length; j++) {
          newarr[j].draw(color(0,255,0));
          this.planeWaves[ind].atomsHit.add(newarr[j]);
          this.sphericalWaves.push(new SphericalWave(newarr[j].x, newarr[j].y));
        }
      }
    }

    for (let ind = 0; ind < this.atoms.length; ind++) {
      if (!this.atoms[ind].wasDrawn) {
        this.atoms[ind].draw();
      }
      this.atoms[ind].wasDrawn = false;
    }

    for (let ind = 0; ind < this.sphericalWaves.length; ind++) {
      if (!this.paused)
        this.sphericalWaves[ind].update();
      if (this.sphericalWaves[ind].cullable()) {
        this.sphericalWaves.splice(ind--, 1);
        continue;
      } else {
        this.sphericalWaves[ind].draw();
      }
    }
  }
}

var sim;
function setup() {

  frameRate(30);
  createCanvas(800, 600);
  fill(0);
  sim = new Simulation();
  sim.bragg_setup();
}

function draw() {
  background(200);
  sim.update();
}