SecondMomentOfInertia.swift

import Foundation
import GIFModule
import HDL
import MM4
import MolecularRenderer
import QuaternionModule
import xTB

// MARK: - User-Facing Options

let renderingOffline: Bool = false
let gifFrameSkipRate: Int = 1

// Frames are recorded and played back at 60 FPS.
let totalSimulationTimePs: Double = 200
let frameSimulationTimePs: Double = 15.0 / 60
let frameCount = Int(totalSimulationTimePs / frameSimulationTimePs)

// 1 degree off works immediately
// 0.5 degrees off requires ~1 full rotation
// 0.1 degrees off takes too long for my patience
let instabilityAngleDegrees: Float = 0.5

// MARK: - Compile Structure

func passivate(topology: inout Topology) {
  func createHydrogen(
    atomID: UInt32,
    orbital: SIMD3<Float>
  ) -> Atom {
    let atom = topology.atoms[Int(atomID)]
    
    var bondLength = atom.element.covalentRadius
    bondLength += Element.hydrogen.covalentRadius
    
    let position = atom.position + bondLength * orbital
    return Atom(position: position, element: .hydrogen)
  }
  
  let orbitalLists = topology.nonbondingOrbitals()
  
  var insertedAtoms: [Atom] = []
  var insertedBonds: [SIMD2<UInt32>] = []
  for atomID in topology.atoms.indices {
    let orbitalList = orbitalLists[atomID]
    for orbital in orbitalList {
      let hydrogen = createHydrogen(
        atomID: UInt32(atomID),
        orbital: orbital)
      let hydrogenID = topology.atoms.count + insertedAtoms.count
      insertedAtoms.append(hydrogen)
      
      let bond = SIMD2(
        UInt32(atomID),
        UInt32(hydrogenID))
      insertedBonds.append(bond)
    }
  }
  topology.atoms += insertedAtoms
  topology.bonds += insertedBonds
}

func createTopology() -> Topology {
  let lattice = Lattice<Hexagonal> { h, k, l in
    let h2k = h + 2 * k
    Bounds { 10 * h + 12 * h2k + 3 * l }
    Material { .checkerboard(.silicon, .carbon) }
  }
  
  var reconstruction = Reconstruction()
  reconstruction.atoms = lattice.atoms
  reconstruction.material = .checkerboard(.silicon, .carbon)
  var topology = reconstruction.compile()
  passivate(topology: &topology)
  return topology
}

func analyze(topology: Topology) {
  print()
  print("atom count:", topology.atoms.count)
  print("hydrogen count:", topology.atoms.count(where: {
    $0.element == .hydrogen
  }))
  do {
    var minimum = SIMD3<Float>(repeating: .greatestFiniteMagnitude)
    var maximum = SIMD3<Float>(repeating: -.greatestFiniteMagnitude)
    for atom in topology.atoms {
      let position = atom.position
      minimum.replace(with: position, where: position .< minimum)
      maximum.replace(with: position, where: position .> maximum)
    }
    print("minimum:", minimum)
    print("maximum:", maximum)
  }
}

var topology = createTopology()
analyze(topology: topology)

// MARK: - Set Up Rotation

func createParameters(
  topology: Topology
) -> MM4Parameters {
  var paramsDesc = MM4ParametersDescriptor()
  paramsDesc.atomicNumbers = topology.atoms.map(\.atomicNumber)
  paramsDesc.bonds = topology.bonds
  return try! MM4Parameters(descriptor: paramsDesc)
}
let parameters = createParameters(topology: topology)

func createRigidBody(
  topology: Topology,
  velocities: [SIMD3<Float>]? = nil
) -> (MM4RigidBody) {
  var rigidBodyDesc = MM4RigidBodyDescriptor()
  rigidBodyDesc.masses = parameters.atoms.masses
  rigidBodyDesc.positions = topology.atoms.map(\.position)
  rigidBodyDesc.velocities = velocities
  let rigidBody = try! MM4RigidBody(descriptor: rigidBodyDesc)
  
  return rigidBody
}

func rotate(topology: inout Topology) {
  let rotation = Quaternion<Float>(
    angle: Float.pi / 180 * instabilityAngleDegrees,
    axis: SIMD3(0, 1, 0))
  let basis0 = rotation.act(on: SIMD3(1, 0, 0))
  let basis1 = rotation.act(on: SIMD3(0, 1, 0))
  let basis2 = rotation.act(on: SIMD3(0, 0, 1))
  
  let rigidBody = createRigidBody(topology: topology)
  let centerOfMass = SIMD3<Float>(rigidBody.centerOfMass)
  
  for atomID in topology.atoms.indices {
    var atom = topology.atoms[atomID]
    let delta = atom.position - centerOfMass
    
    var rotatedDelta: SIMD3<Float> = .zero
    rotatedDelta += basis0 * delta[0]
    rotatedDelta += basis1 * delta[1]
    rotatedDelta += basis2 * delta[2]
    
    atom.position = centerOfMass + rotatedDelta
    topology.atoms[atomID] = atom
  }
}

// Rough estimate of linear velocity from rotation. Just to make sure the
// starting value is not absurdly fast.
//
// 6.45 nm Y dimension -> 3.23 nm radius
// v = ωr
// v = 2πfr
//
// v = (2π)(1 GHz)(3.23 nm)
// v = 20 m/s per GHz
//
// 1 GHz would require 1000 ps of simulation to see one rotation, which
// is pretty slow. 20 m/s is also far too slow for reasonably quick MD tests.
//
// Velocity of gas molecules at room temperature:
// H2: 1754 m/s
// He: 1245 m/s
// N2: 470-515 m/s
// CO2: 375 m/s
//
// 30 GHz -> 600 m/s -> wait for 33 ps of simulation to visualize 1 rotation
//
// None of this kinetic energy should get thermalized because it's a bulk DOF
// of the rigid body itself, not distributed among internal vibrations
// between the atoms. Conservation of momentum means the bulk DOF cannot get
// redistributed, unless it collides with another nanomachine.
func rotatedVelocities(topology: Topology) -> [SIMD3<Float>] {
  let frequencyGHz = SIMD3<Float>(30, 0, 0) // x-axis approximately 2nd moment
  let angularVelocityRadNs = frequencyGHz * 2 * Float.pi
  let angularVelocityRadPs = angularVelocityRadNs / 1000
  
  let rigidBody = createRigidBody(topology: topology)
  let centerOfMass = SIMD3<Float>(rigidBody.centerOfMass)
  
  var output: [SIMD3<Float>] = []
  for atomID in topology.atoms.indices {
    var atom = topology.atoms[atomID]
    let radius = atom.position - centerOfMass
    
    // v = ω x r
    func cross(
      _ lhs: SIMD3<Float>,
      _ rhs: SIMD3<Float>
    ) -> SIMD3<Float> {
      let yzx = SIMD3<Int>(1, 2, 0)
      let zxy = SIMD3<Int>(2, 0, 1)
      return (lhs[yzx] * rhs[zxy]) - (lhs[zxy] * rhs[yzx])
    }
    let linearVelocity = cross(angularVelocityRadPs, radius)
    output.append(linearVelocity)
  }
  return output
}

rotate(topology: &topology)
let velocities = rotatedVelocities(topology: topology)
let rigidBody = createRigidBody(
  topology: topology,
  velocities: velocities)

// We'll rotate about the 2nd principal axis, the unstable one.
print(SIMD3<Float>(rigidBody.principalAxes.0)) // 1st axis
print(SIMD3<Float>(rigidBody.principalAxes.1)) // 2nd axis
print(SIMD3<Float>(rigidBody.principalAxes.2)) // 3rd axis
print("moment of inertia:", SIMD3<Float>(rigidBody.momentOfInertia))

// Double check that we set the rotation frequency correctly.
do {
  // rad/ps -> rad/ns -> GHz
  let angularVelocity = rigidBody.angularMomentum / rigidBody.momentOfInertia
  let angularVelocityRadNs = angularVelocity * 1000
  let frequencyGHz = angularVelocityRadNs / (2 * Double.pi)
  print("frequency:", SIMD3<Float>(frequencyGHz), "GHz")
}

// Check that the maximum velocity of any atom is indeed ~600 m/s.
do {
  var maxSpeed: Float = .zero
  for atomID in topology.atoms.indices {
    let velocity = rigidBody.velocities[atomID]
    let speed = (velocity * velocity).sum().squareRoot()
    if speed > maxSpeed {
      maxSpeed = speed
    }
  }
  print("max speed of individual atom:", maxSpeed * 1000, "m/s")
}

// MARK: - Run Simulation

@MainActor
func createForceField() -> MM4ForceField {
  var forceFieldParameters = MM4Parameters()
  forceFieldParameters.append(contentsOf: parameters)
  
  var forceFieldDesc = MM4ForceFieldDescriptor()
  forceFieldDesc.integrator = .verlet
  forceFieldDesc.parameters = forceFieldParameters
  let forceField = try! MM4ForceField(descriptor: forceFieldDesc)
  
  forceField.positions = rigidBody.positions
  forceField.velocities = rigidBody.velocities
  return forceField
}
let forceField = createForceField()

var frames: [[Atom]] = []
@MainActor
func createFrame(positions: [SIMD3<Float>]) -> [Atom] {
  var output: [SIMD4<Float>] = []
  for atomID in topology.atoms.indices {
    var atom = topology.atoms[atomID]
    atom.position = positions[atomID]
    output.append(atom)
  }
  return output
}

for frameID in 1...frameCount {
  forceField.simulate(time: frameSimulationTimePs)
  
  var rigidBodies: [MM4RigidBody] = []
  for rigidBodyID in 0..<1 {
    var positions: [SIMD3<Float>] = []
    var velocities: [SIMD3<Float>] = []
    
    let baseAddress = rigidBodyID * topology.atoms.count
    for atomID in topology.atoms.indices {
      let position = forceField.positions[baseAddress + atomID]
      let velocity = forceField.velocities[baseAddress + atomID]
      positions.append(position)
      velocities.append(velocity)
    }
    
    var rigidBodyDesc = MM4RigidBodyDescriptor()
    rigidBodyDesc.masses = parameters.atoms.masses
    rigidBodyDesc.positions = positions
    rigidBodyDesc.velocities = velocities
    let rigidBody = try! MM4RigidBody(descriptor: rigidBodyDesc)
    rigidBodies.append(rigidBody)
  }
  
  let time = Double(frameID) * frameSimulationTimePs
  print()
  print("t = \(String(format: "%.3f", time)) ps")
  
  for rigidBodyID in 0..<1 {
    // Changes from (0, 30.0, 0) at the start to (24.2, 4.8, 25.0) at one point.
    let rigidBody = rigidBodies[rigidBodyID]
    let angularVelocity = rigidBody.angularMomentum / rigidBody.momentOfInertia
    let angularVelocityRadNs = angularVelocity * 1000
    let frequencyGHz = angularVelocityRadNs / (2 * Double.pi)
    print("- frequency:", SIMD3<Float>(frequencyGHz), "GHz")
    
    // Always stays within the envelope of 6500-6700 zJ.
    let angularEnergy =
    0.5 * (angularVelocity * rigidBody.angularMomentum).sum()
    let angularEnergyRepr = String(format: "%.1f", angularEnergy)
    print("- angular kinetic energy:", angularEnergyRepr, "zJ")
  }
  
  let positions = forceField.positions
  let frame = createFrame(positions: positions)
  frames.append(frame)
}

// Input: time in seconds
// Output: atoms
func interpolate(
  frames: [[Atom]],
  time: Float
) -> [Atom] {
  guard frames.count >= 1 else {
    fatalError("Need at least one frame to know size of atom list.")
  }
  
  let multiple60Hz = time * 60
  var lowFrame = Int(multiple60Hz.rounded(.down))
  var highFrame = lowFrame + 1
  var lowInterpolationFactor = Float(highFrame) - multiple60Hz
  var highInterpolationFactor = multiple60Hz - Float(lowFrame)
  
  if lowFrame < -1 {
    fatalError("This should never happen.")
  }
  if lowFrame >= frames.count - 1 {
    lowFrame = frames.count - 1
    highFrame = frames.count - 1
    lowInterpolationFactor = 1
    highInterpolationFactor = 0
  }
  
  var output: [Atom] = []
  for atomID in frames[0].indices {
    let lowAtom = frames[lowFrame][atomID]
    let highAtom = frames[highFrame][atomID]
    
    var position: SIMD3<Float> = .zero
    position += lowAtom.position * lowInterpolationFactor
    position += highAtom.position * highInterpolationFactor
    
    var atom = lowAtom
    atom.position = position
    output.append(atom)
  }
  return output
}

// MARK: - Launch Application

@MainActor
func createApplication() -> Application {
  // Set up the device.
  var deviceDesc = DeviceDescriptor()
  deviceDesc.deviceID = Device.fastestDeviceID
  let device = Device(descriptor: deviceDesc)
  
  // Set up the display.
  var displayDesc = DisplayDescriptor()
  displayDesc.device = device
  if renderingOffline {
    displayDesc.frameBufferSize = SIMD2<Int>(1080, 1080)
  } else {
    displayDesc.frameBufferSize = SIMD2<Int>(1440, 1440)
  }
  if !renderingOffline {
    displayDesc.monitorID = device.fastestMonitorID
  }
  let display = Display(descriptor: displayDesc)
  
  // Set up the application.
  var applicationDesc = ApplicationDescriptor()
  applicationDesc.device = device
  applicationDesc.display = display
  if renderingOffline {
    applicationDesc.upscaleFactor = 1
  } else {
    applicationDesc.upscaleFactor = 3
  }
  
  applicationDesc.addressSpaceSize = 4_000_000
  applicationDesc.voxelAllocationSize = 500_000_000
  applicationDesc.worldDimension = 64
  let application = Application(descriptor: applicationDesc)
  
  return application
}
let application = createApplication()

for atomID in topology.atoms.indices {
  let atom = topology.atoms[atomID]
  application.atoms[atomID] = atom
}

@MainActor
func createTime() -> Float {
  if renderingOffline {
    let elapsedFrames = max(0, gifFrameSkipRate * application.frameID - 30)
    let frameRate: Int = 60
    let seconds = Float(elapsedFrames) / Float(frameRate)
    return seconds
  } else {
    let elapsedFrames = application.clock.frames
    let frameRate = application.display.frameRate
    let seconds = Float(elapsedFrames) / Float(frameRate)
    return seconds
  }
}

@MainActor
func modifyAtoms() {
  let time = createTime()
  let atoms = interpolate(frames: frames, time: time)
  for atomID in atoms.indices {
    let atom = atoms[atomID]
    application.atoms[atomID] = atom
  }
}

@MainActor
func modifyCamera() {
  let rotation = Quaternion<Float>(
    angle: Float.pi / 180 * 0,
    axis: SIMD3(1, 0, 0))
  
  func transform(_ input: SIMD3<Float>) -> SIMD3<Float> {
    var output = input
    output = rotation.act(on: output)
    return output
  }
  
  application.camera.basis.0 = transform(SIMD3(1, 0, 0))
  application.camera.basis.1 = transform(SIMD3(0, 1, 0))
  application.camera.basis.2 = transform(SIMD3(0, 0, 1))
  
  application.camera.position = transform(SIMD3(1.3, 2.8, 10))
}

// Enter the run loop.
if !renderingOffline {
  application.run {
    modifyAtoms()
    modifyCamera()
    
    var image = application.render()
    image = application.upscale(image: image)
    application.present(image: image)
  }
} else {
  let frameBufferSize = application.display.frameBufferSize
  var gif = GIF(
    width: frameBufferSize[0],
    height: frameBufferSize[1])
  
  print("rendering frames")
  for _ in 0..<((30 + frameCount) / gifFrameSkipRate) {
    let loopStartCheckpoint = Date()
    modifyAtoms()
    modifyCamera()
    
    let image = application.render()
    
    var gifImage = GIFModule.Image(
      width: frameBufferSize[0],
      height: frameBufferSize[1])
    for y in 0..<frameBufferSize[1] {
      for x in 0..<frameBufferSize[0] {
        let address = y * frameBufferSize[0] + x
        
        // Leaving this in the original SIMD4<Float16> causes a CPU-side
        // bottleneck on Windows.
        let pixel = SIMD4<Float>(image.pixels[address])
        
        // Don't clamp to [0, 255] range to avoid a minor CPU-side bottleneck.
        // It theoretically should never go outside this range; we just lose
        // the ability to assert this.
        let scaled = pixel * 255
        
        // On the Windows machine, '.toNearestOrEven' causes a massive
        // CPU-side bottleneck.
        let rounded = (scaled + 0.5).rounded(.down)
        
        // Avoid massive CPU-side bottleneck for unknown reason when casting
        // floating point vector to integer vector.
        let r = UInt8(rounded[0])
        let g = UInt8(rounded[1])
        let b = UInt8(rounded[2])
        
        let color = Color(
          red: r,
          green: g,
          blue: b)
        
        gifImage[y, x] = color
      }
    }
    
    let quantization = OctreeQuantization(fromImage: gifImage)
    
    let frame = Frame(
      image: gifImage,
      delayTime: 4, // 25 FPS
      localQuantization: quantization)
    gif.frames.append(frame)
    
    let loopEndCheckpoint = Date()
    print(loopEndCheckpoint.timeIntervalSince(loopStartCheckpoint))
  }
  
  print("encoding GIF")
  let encodeStartCheckpoint = Date()
  let data = try! gif.encoded()
  let encodeEndCheckpoint = Date()
  
  let encodedSizeRepr = String(format: "%.1f", Float(data.count) / 1e6)
  print("encoded size:", encodedSizeRepr, "MB")
  print(encodeEndCheckpoint.timeIntervalSince(encodeStartCheckpoint))
  
  let packagePath = FileManager.default.currentDirectoryPath
  let filePath = "\(packagePath)/.build/video.gif"
  let succeeded = FileManager.default.createFile(
    atPath: filePath,
    contents: data)
  guard succeeded else {
    fatalError("Could not write to file.")
  }
}