// Calculation of the movement of the launch loop during launch
// By Kristian Andresen
// To do:
// - Conservation of energy
// - Planetary rotation
// - Pressure as function of height
// - Wind loading
// - Spaceship launch

#include <iostream>
#include <fstream>
#include <stdlib.h>
#include <math.h>
#include <sstream> 
#include <vector>
using namespace std;

// Utilities
namespace transient {
#define swap(T,a,b) { T buff = b; b = a; a = buff; }
double temp[3]; // Temporary holder of vectors
#define plusMinus(i) for(int i = -1; i <= 1; i = i+2)
#define for2(i) for(int i = 0; i < 2; i++)
#define for3(i) for(int i = 0; i < 3; i++)
#define forN(N,i) for(int i = 0; i < N; i++)
inline double sqr(double x) { return x*x; }
inline double dot(double a[], double b[]) { return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]; }
inline double sqr(double a[]) { return dot(a, a); }
inline double* sub(double res[], double a[], double b[]) { for3(i) res[i] = a[i] - b[i]; return res; }
inline double* sub(double a[], double b[]) { for3(i) temp[i] = a[i] - b[i]; return temp; }
inline void normalize(double a[]) { double length = sqrt(sqr(a)); for3(i) a[i] /= length; } 
inline double dist(double a[], double b[]) { return sqrt(sqr(sub(a, b))); }
inline int wrap(int i) { return i >= 3? i-3 : i; }
inline double abs(double a) { return a < 0? -a : a; } 
inline double* cross(double c[], double a[], double b[]) { for3(i) c[i] = a[wrap(1+i)]*b[wrap(2+i)] - a[wrap(2+i)]*b[wrap(1+i)]; return c; }
int getColor(int r, int g, int b) { return r * 65536 + g * 256 + b; }
const int rotorColor[4] = {getColor(255,0,0), getColor(0,255,0), getColor(255,255,0), getColor(0,255,255)};
}

// Constants and globals
namespace transient {
// Physical constants
const double earthRadius = 6378100; // [m]
const double earthMassTimesG = 3.98734836e14; // [N*m^2/kg]
const double earthMass = 5.9736e24; // [kg]
const double pi = 3.1415926;

// Track design parameters
const double distance = 100e3; // [m] Distance between ground stations
const double rotorDensity = 1; // [kg/m] Rotor density at apex
const double apex[3] = {0.0, 0.0, earthRadius + 40e3}; // [x,y,z] 
const double setupApexSpeed = 2000; // [m/s] Speed of rotor at apex during initialization
const double simulationApexSpeed = 2000; // [m/s] Speed of rotor at apex during simulation
const double soundInTrack = 4000; // [m/s] Speed of sound in track
const double specificStrength = 1e5; // [N/(kg/m)] Maximum tensile force in track depends on density
const double trackArea = 0.001; // [m^2] This influences track resistance to bending
const double shearModulus = 60e9; // [Pa] Shear modulus of steel
const double youngModulus = 200e9; // [Pa] Youngs modulus of steel
const int tracks = 4; // Number of tracks
const bool eastbound[4] = {true, false, true, false}; // Direction of each track
const bool helicalTrack = true; // If true, track is initialized as a spiral
const int guyWires = 5; // Number of guy wires attached to each track segment, uneven number for balance
const bool guyWire = (tracks > 1) && helicalTrack; // If true, employ guy wires
const double helixRadius = 1000; // [m]
const double helixes = 7; // Number of twists along the track, need not be integer
const double centrifugalPeriod = distance * 2 / simulationApexSpeed / helixes; // [s]
const double centrifugalAcceleration = 4 * sqr(pi) * helixRadius / sqr(centrifugalPeriod); // [m/s^2]

// Simulation parameters
const int segments = 200; // Defines the resolution of the calculation, even number
const int timesteps = 20000; // Length of simulation
const int sampling = 200; // Data is sampled after this number of timesteps
const double setupTick = distance / segments / setupApexSpeed; // [s] Duration of one tick during initialization
// If simulationTick is too large, simulation will be numerically unstable, but not for reasons which have to do with physics
const double simulationTick = setupTick / 100; // [s] A method is needed to calculate a suitable simulationTick

// Globals
double tick; // [s] Duration of one time step
double apexSpeed; // [m/s] Speed of rotor at apex
double trackDensity; // [kg/m]
}

namespace transient {

struct Spring; // Forward declaration

// Main data structure
struct Segment
{
  // Segment attributes
  double position[3]; // [m]
  double velocity[3]; // [m/s]
  double rotorMass; // [kg]
  double trackMass; // [kg] 
  double massRatio; // Ratio between gravitational and inertial mass, used to apply gravity during setup
  double direction[3]; // Normalized vector to next track segment

  // Mass and momentum transfer variables
  double designSpeed; // [m/s]
  double rotorSpeed; // [m/s] Actual speed
  double stableSpeed; // [m/s]
  double incomingMass; // [kg]
  double incomingVelocity[3]; // [m/s]
  double incomingPosition[3]; // [m]

  // Spring variables
  Spring* shearStress[2][2]; // Springs to model bending stresses [0 is upstream, 1 is downstram][0 is down, 1 is sideways]
  Spring* track; // Connection to downstream track segment
  Spring* wire[tracks][guyWires]; // Guy wires

  Segment() { 
    for2(i) for2(j) shearStress[i][j] = 0;
    track = 0;
    forN(tracks,d) forN(guyWires,w) wire[d][w] = 0;
  };
  double mass() { return rotorMass + trackMass; };
  double length() { return dist(position, (this+1)->position); };
  double speed() { return sqrt(sqr(velocity)); };
  double height() { return sqrt(sqr(position)) - earthRadius; }; // Height above ground
  void findDirection() { normalize(sub(direction, (this+1)->position, position)); };
  void applyGravity();
  void handleMomentum() { for3(a) position[a] += velocity[a] * tick; };
  void handleArrival();
  void activeStabilization();
  void handleDeparture();
  void connect(Spring** c, Segment* s, double* vector, double length, bool setup);
  void connectTrack(Segment* s, bool setup);
  void connectWire(Segment* s, int i, int j, int w, bool setup);
  void shearForces(bool setup);
  double* guyForce();
};

// Data structure for guy wires, track connections and shear stress, modelled as Kelvin-Voigt materials
struct Spring // Each spring has design length, present length, and old length from previous timestep
{
  bool bendable; // True for guy wires, false for track
  double designLength; // [m]
  double length; // [m] The length can be negative when modelling shear forces with springs
  double oldLength; // [m] Used to find damping
  double designTension; // [N]
  double guyTension; // [N] Only used by guy wires
  Segment* end[2]; // Remember that pointers can't handle being copied
  double dampingForce; // [N]
  double springConstant; // [kg/s^2]
  double vector[3]; // [m] Projection of vector from end[0] to end[1] onto spring direction
  void contract();
  double dampingConstant() { return 2 * sqrt(springConstant * (end[0]->mass() + end[1]->mass())) / 100; }; // [kg/s] Guess
  double springForce() { return - springConstant * (length - designLength) - designTension; };
  Spring() { guyTension = 0; };
};

// Data
Segment track[tracks][segments+1];
vector<Spring*> spring;

// Member functions
void Segment::applyGravity() {
  double down[3];
  for3(a) down[a] = -position[a];
  double g = earthMassTimesG * massRatio / sqr(down);
  normalize(down);
  for3(a) velocity[a] += down[a] * g * tick;
  for3(a) position[a] += down[a] * g * sqr(tick) / 2;
}
void Segment::handleArrival() {
  for3(a) velocity[a] = (incomingMass * incomingVelocity[a] + mass() * velocity[a]) / (mass() + incomingMass); 
  for3(a) position[a] = (incomingMass * incomingPosition[a] + mass() * position[a]) / (mass() + incomingMass);
  rotorMass += incomingMass;
}
void Segment::activeStabilization() { // This doesn't perform much better than simply setting stableSpeed = designSpeed
  stableSpeed = mass() * dot(velocity, direction) / rotorMass / sqr(direction); // Minimizing kinetic energy of track
  double minMaxSpeed[2]; // Index 0 is not necessarily the minimum
  for2(i) { // Alternatively, choose rotor speed such that potential plus kinetic energy of rotor is constant
    double g = earthMassTimesG / sqr((this-1+i)->position);
    minMaxSpeed[i] = sqrt(sqr((this-1+2*i)->rotorSpeed)
      - 2 * (this-1+i)->mass() / (this-1+i)->rotorMass * g * (this->height() - (this-1+2*i)->height()));
  }
  if ((stableSpeed > minMaxSpeed[0] && stableSpeed < minMaxSpeed[1]) || (stableSpeed < minMaxSpeed[0] && stableSpeed > minMaxSpeed[1])) {

  } else stableSpeed = designSpeed;
}
void Segment::handleDeparture() { 
  rotorSpeed = stableSpeed;
  // First we find the mass to be moved based on the speed and density
  (this+1)->incomingMass = rotorSpeed * rotorMass / length() * tick;
  rotorMass -= (this+1)->incomingMass;
  // Next, find the outgoing velocity
  normalize(sub((this+1)->incomingVelocity, (this+1)->position, position));
  for3(a) (this+1)->incomingVelocity[a] *= rotorSpeed;
  // Next, conserve momentum
  for3(a) velocity[a] = ((mass() + (this+1)->incomingMass) * velocity[a] - (this+1)->incomingMass * (this+1)->incomingVelocity[a]) / mass();
  // Lastly, adjust center of mass
  for3(a) (this+1)->incomingPosition[a] = (position[a] + (this+1)->position[a]) / 2;
  for3(a) position[a] = ((mass() + (this+1)->incomingMass) * position[a] - (this+1)->incomingMass * (this+1)->incomingPosition[a]) / mass();
}
void Segment::connect(Spring** c, Segment* s, double* vector, double length, bool setup) {
  if (setup) {
    spring.push_back((*c) = new Spring());
    (*c)->designLength = length;
    (*c)->oldLength = length;
    (*c)->end[0] = this;
    (*c)->end[1] = s;
    (*c)->bendable = false;
    (*c)->designTension = 0;
  }
  for3(a) (*c)->vector[a] = vector[a];
  (*c)->length = length;
}
void Segment::connectTrack(Segment* s, bool setup) {
  Spring** c = &track;
  connect(c, s, sub(s->position, position), dist(s->position, position), setup);
  if (setup) {
    (*c)->springConstant = sqr(soundInTrack) * trackDensity / (*c)->length;
    (*c)->designTension = 1e4; 
  } else (*c)->contract();
}
void Segment::connectWire(Segment* s, int i, int j, int w, bool setup) {
  Spring** c = &wire[j][w];
  connect(c, s, sub(s->position, position), dist(s->position, position), setup);
  if (setup) {
    (*c)->springConstant = sqr(soundInTrack) * trackDensity / (*c)->length / 10; // Artificially high, active tensioning
    (*c)->bendable = true;
    (*c)->designTension = 1; // Temporary value
    s->wire[i][w] = (*c); // Give opposing segment access to spring data
  } else (*c)->contract();
}
void Segment::shearForces(bool setup) {
  // First, find guy wire and track direction - these change with time and cannot be cached
  double axis[3][3]; // [0 is the track direction, 1 is down, 2 is sideways][x,y,z]
  for3(i) for3(j) axis[i][j] = 0;
  sub(axis[0], (this+1)->position, (this-1)->position);
  if (guyWire) { // Use guy wires to define the local "down" direction - should perhaps also use neighbour segments
    forN(tracks,d) forN(guyWires,w) if (wire[d][w]) for3(a) axis[1][a] += (wire[d][w]->end[0] == this? 1 : -1) * wire[d][w]->vector[a];
    cross(axis[2], axis[0], axis[1]);
  } else { // Use north to define sideways
    axis[2][1] = 1;
    cross(axis[1], axis[0], axis[2]);
    cross(axis[2], axis[0], axis[1]);
  }
  for3(a) normalize(axis[a]);
  // Next, do projections, reducing to one dimensional strains
  double projection[3][3]; // Indices are segment and axis 
  for3(s) for3(a) projection[s][a] = dot((this-1+s)->position, axis[a]) / sqr(axis[a]);
  // Next, find the deflection
  double deflection[2]; // [0 is guy wire direction, 1 is sideways]
  for2(i) deflection[i] = projection[1][i+1] - (projection[0][i+1] + projection[2][i+1])/2;
  // Move the data to the spring
  for2(i) for2(j) {
    Spring** c = &shearStress[i][j];
    connect(c, this+2*i-1, axis[1+j], -deflection[j], setup);
    if (setup) (*c)->springConstant = shearModulus * trackArea / length();
    else (*c)->contract();
  }
}
double* Segment::guyForce() { // Find the combined force of all guy wires
  for3(a) temp[a] = 0;
  forN(tracks,j) forN(guyWires,w) if (wire[j][w]) {
    normalize(wire[j][w]->vector);
    for3(a) temp[a] += (wire[j][w]->end[0] == this? -1 : 1) * wire[j][w]->springForce() * wire[j][w]->vector[a];
  }
  return temp;
}

void Spring::contract() {
  double springSpeed = (length - oldLength) / tick;
  oldLength = length;
  dampingForce = - dampingConstant() * springSpeed;
  normalize(vector); // Get direction of acceleration
  double force = springForce() + dampingForce;
  if (bendable && force > 0) force = 0; // A bendable spring can only provide a contracting (negative) force, or it buckles
  for2(i) for3(a) end[i]->velocity[a] += (2*i-1) * force * tick * vector[a] / end[i]->mass(); // Newtons law
}

// Initialization
void traceTrack(Segment* t) { 
  // Define initial parabola - disperse track segments so the slugs pass by them at constant time intervals
  forN(segments+1,s) t[s].rotorMass = rotorDensity * apexSpeed * tick; // Is this balanced right?
  Segment* middle = &t[segments/2];
  for3(a) middle->position[a] = apex[a];
  middle->velocity[0] = apexSpeed;
  middle->trackMass = trackDensity * apexSpeed * tick;
  middle->massRatio = (middle->trackMass + middle->rotorMass) / middle->rotorMass;
  for (int d = -1; d <= 1; d += 2) { // 2 different directions
    for3(a) middle->velocity[a] *= -1; // Switch direction
    for (int i = 1; i <= segments/2; i++) { 
      Segment* s = t + segments/2 + d*i;
      for3(a) s->position[a] = (s-d)->position[a] + (s-d)->velocity[a] * tick;
      for3(a) s->velocity[a] = (s-d)->velocity[a];
      s->trackMass = trackDensity * s->speed() * tick;
      s->massRatio = s->mass() / s->rotorMass;
      s->applyGravity();
    }
  }
  forN(segments+1,s) t[s].position[1] = (double(rand())/RAND_MAX-0.5)/100; // Centimeter precision in startup
}
void findTrackDensity(Segment* t) {
  // Adjust the weight of the track to keep the rotor down
  trackDensity = 10 * rotorDensity; // Guess, to be improved upon
  for (double adjustment = trackDensity/2; adjustment > 0.01; adjustment /= 2) {
    traceTrack(t);
    trackDensity += adjustment * (t[0].height() > 0? 1 : -1);
  }
  cout << "Track density " << trackDensity << endl;
  cout << "Rotor speed at ground " << t[0].speed() << endl;
}
void applyHelix(Segment* t, int d)
{
  double north[3] = {0,1,0};
  double up[segments+1][3];
  forN(segments+1,s) {
    normalize(cross(up[s], t[s].direction, north));
    // Equal time between segments leads to equal periods of centrifugal rotation
    for3(a) t[s].position[a] += north[a] * helixRadius * sin(helixes*pi*s/segments + d*2*pi/tracks);
    for3(a) t[s].position[a] += up[s][a] * helixRadius * cos(helixes*pi*s/segments + d*2*pi/tracks);
  }
}
void handleSprings(bool setup) {
  if (guyWire) forN(segments+1,i) forN(tracks,d) forN(tracks,j) forN(guyWires,w) {
    int opposite = i + w - (guyWires-1)/2;
    if (eastbound[d] != eastbound[j]) opposite = segments - opposite; 
    Segment* s = &track[d][i];
    if (opposite >= 0 && opposite <= segments && d != j) { // Stay within track, don't connect to self
      if ((setup && !s->wire[j][w] && !track[j][opposite].wire[d][w]) // Check that this isn't a repeat
        || (!setup && s->wire[j][w] && s->wire[j][w]->end[0] == s)) { // Check that wire exists
        s->connectWire(&track[j][opposite], d, j, w, setup);
      }
    }
  }
  // Fix guy wire tension so it balances centrifugal force
  if (setup && guyWire) forN(segments+1,i) forN(tracks,d) {
    Segment* s = &track[d][i];
    double* sum = s->guyForce();
    double factor = centrifugalAcceleration * s->rotorMass / s->length() * s->designSpeed * tick / 2 / sqrt(sqr(sum));
    forN(tracks,j) forN(guyWires,w) if (s->wire[j][w]) s->wire[j][w]->guyTension += factor;
  }
  if (setup && guyWire) forN(segments+1,s) forN(tracks,i) forN(tracks,j) forN(guyWires,w)
    if (track[i][s].wire[j][w]) track[i][s].wire[j][w]->designTension = track[i][s].wire[j][w]->guyTension;
  // Track connection
  forN(tracks,d) forN(segments,s) track[d][s].connectTrack(&track[d][s+1], setup);
  // Shear forces
  forN(tracks,d) forN(segments-1,s) track[d][s+1].shearForces(setup);
}
void setupLoop() {
  findTrackDensity(track[0]);
  forN(tracks,d) forN(segments+1,s) track[d][s] = track[0][s]; // Copy initial parabola
  forN(tracks,d) {
    forN(segments,s) track[d][s].findDirection();
    for3(a) track[d][segments].direction[a] = track[d][segments-1].direction[a]; // This is an approximation
    if (helicalTrack) applyHelix(track[d], d);
    // The westbound track must move opposite the eastbound track, so we swap
    if (!eastbound[d]) forN(segments/2,s) swap(Segment,track[d][s],track[d][segments-s]); 
    forN(segments,s) track[d][s].findDirection(); // Direction has changed now
    forN(segments+1,s) track[d][s].designSpeed = track[d][s].speed() * (simulationApexSpeed/setupApexSpeed);
    forN(segments+1,s) track[d][s].rotorSpeed = track[d][s].designSpeed; 
    forN(segments+1,s) for3(a) track[d][s].velocity[a] = track[d][s].designSpeed * track[d][s].direction[a] / track[d][s].massRatio;
    forN(segments+1,s) track[d][s].massRatio = 1; 
  }
  handleSprings(true); // Setup of springs
}

// Output
void saveToFile(int t) {
  ofstream file;
  stringstream filename;
  filename << "data/t" << 1000000 + t << ".txt"; // Inelegant method for getting alphabetical order
  file.open(filename.str());
  forN(tracks,d) forN(segments-1,i)  {
    Segment* s = &track[d][i+1];
    int color = (i + int((timesteps-t) * (simulationApexSpeed/setupApexSpeed) / (setupTick/simulationTick))) % (segments/2);
    if (color > segments/4) color = rotorColor[d]; else color = getColor(0,0,0);
    if (i == 0) color = getColor(255,255,255); // Invisible jump from eastbound to westbound
    file << s->position[0] / 1000 << " ";
    file << s->position[1] << " ";
    file << s->height() / 1000 << " ";
    file << color << " ";
    file << log(abs(s->track->springForce())+10e-20) << " ";
    file << log(sqrt(sqr(s->guyForce()))+10e-20) << " ";
    file << log(s->rotorMass/track[d][i].length()*sqr(s->designSpeed)+10e-20) << "\n";
  }
  file.close(); 
}
void evaluate() {
  // Rotor mass
  double maxRotor[4];
  for2(i) maxRotor[i] = track[0][0].rotorMass; for2(i) maxRotor[2+i] = 0;
  forN(tracks,t) forN(segments,s) for2(i) if (track[t][s].rotorMass * (2*i-1) > (2*i-1) * maxRotor[i]) {
    maxRotor[i] = track[t][s].rotorMass;
    maxRotor[2+i] = s;
  }
  cout << "Rotor mass extremum "; for2(i) cout << 100*maxRotor[i]/track[0][0].rotorMass << "% ";
  cout << "Index "; for2(i) cout << maxRotor[2+i] << " "; cout << endl;
  // Rotor speed
  for2(i) maxRotor[i] = 1; for2(i) maxRotor[2+i] = 0;
  forN(tracks,t) forN(segments-1,s) for2(i) if (track[t][s+1].rotorSpeed/track[t][s+1].designSpeed * (2*i-1) > (2*i-1) * maxRotor[i]) {
    maxRotor[i] = track[t][s+1].rotorSpeed/track[t][s+1].designSpeed;
    maxRotor[2+i] = s+1;
  }
  cout << "Rotor speed extremum "; for2(i) cout << 100*maxRotor[i] << "% ";
  cout << "Index "; for2(i) cout << maxRotor[2+i] << " "; cout << endl;
}

// Simulation
void simulateTrack(Segment* t) {
  for (int i = 1; i < segments; i++) t[i].findDirection();
  // Handle the injection of new slugs
  t[1].incomingMass = t[0].rotorMass / t[0].length() * t[0].designSpeed * tick;
  for3(a) t[1].incomingVelocity[a] = t[0].designSpeed * t[0].direction[a];
  for3(a) t[1].incomingPosition[a] = (t[0].position[a] + t[1].position[a]) / 2;
  // Handle the main track
  for (int i = 1; i < segments; i++) t[i].applyGravity();
  for (int i = 1; i < segments; i++) t[i].activeStabilization();
  for (int i = 1; i < segments; i++) t[i].handleDeparture();
  for (int i = 1; i < segments; i++) t[i].handleArrival();
  for (int i = 1; i < segments; i++) t[i].handleMomentum();
}
void simulateLoop() {
  cout << "Remaining timesteps ";
  forN(timesteps+1,t) {
    if (t % (timesteps/10) == 0) { cout << 100 - 100*t/timesteps << "%" << endl; if (t > 0) evaluate(); }
    else if (t % (timesteps/100) == 0) cout << ".";
    forN(tracks,d) simulateTrack(track[d]);
    handleSprings(false);
    if (t % sampling == 0) saveToFile(t);
  }
}

} // End of namespace

using namespace transient;

int main () {
  // Initialize
  tick = setupTick; 
  apexSpeed = setupApexSpeed;
  setupLoop();
  // Simulate
  tick = simulationTick;
  apexSpeed = simulationApexSpeed;
  simulateLoop();
  return 0;
}