#!/usr/bin/env python
from math import sqrt
from util import readFile, writeFile
import sys
from getopt import getopt
from state import State
import numpy as np
from matplotlib import pyplot as plt
from constants import k_B, m

def main(argv):
  if len(argv) != 4:
    print(f'{argv[0]} <number of particles> <side length> <temperature>')
    sys.exit(2)

  state = State(
    count=int(argv[1]),
    size=float(argv[2]), 
    temperature=float(argv[3])
  )

  state.molecules = np.hstack([np.random.rand(state.count, 3)*state.size, np.zeros((state.count, 3))])

  state.minimize()

  sigma = sqrt(k_B*state.temperature/m)
  state.molecules[:,3:6] = np.random.multivariate_normal([0,0,0], np.eye(3)*sigma, state.count)

  avg = np.sum(state.molecules[:,3:6], 0)
  state.molecules[:,3:6] -= avg / state.count

  writeFile('./res/initial.mol', state)


if __name__ == '__main__':
  main(sys.argv)