# GHSCOM with python

import numpy as np
import matplotlib.pyplot as plt

# use pywin32 to access GHSCOM
import win32com.client
ghs = win32com.client.Dispatch("GHSCOM.Application")

# GHS run file command
Ntrims=9
Trim1=-20
TrimInc=5
Nheels=17
Heel1=-40
HeelInc=5

ghs.Command("clear")

ghs.Command("read box2.gf")
ghs.Command("we 1500 10.714 0 11.25")
ghs.Command("solve")

ghs.Command("macro getBGN")
ghs.Command(" trim %1")
ghs.Command(" macro setHeel")
ghs.Command("  heel %%91")
ghs.Command("  solve depth")
ghs.Command("  variable BGN_T{trim:0}_H{heel:0}={BGN}")
ghs.Command(" //")
ghs.Command(" .setHeel ({Nheels},{HeelInc}) {Heel1}")
ghs.Command("/")
ghs.Command("variable Ntrims, Nheels, Trim1, Heel1, TrimInc, HeelInc")
ghs.Command("set Ntrims="+str(Ntrims))
ghs.Command("set Trim1="+str(Trim1))
ghs.Command("set TrimInc="+str(TrimInc))
ghs.Command("set Nheels="+str(Nheels))
ghs.Command("set Heel1="+str(Heel1))
ghs.Command("set HeelInc="+str(HeelInc))

ghs.Command(".getBGN ({Ntrims},{TrimInc}) {Trim1}")

#=======================================================
# copy GHS variables to python variables
gfname=ghs.GetVariableString("GF")
print('GF: ',gfname)

emptyArray = np.zeros([Nheels,Ntrims])

Trims = np.arange(Trim1,Trim1+TrimInc*Ntrims,TrimInc)
Heels = np.arange(Heel1,Heel1+HeelInc*Nheels,HeelInc)
print(Trims)
print(Heels)

for i in range(Ntrims):
  for j in range(Nheels):
    vname = 'BGN_T{}_H{}'
    GHSvar=ghs.GetVariableString(vname.format(Trims[i],Heels[j]))
    print(vname.format(Trims[i],Heels[j]),"=",GHSvar)
    emptyArray[j,i]=GHSvar

del ghs # kill the COM object

plt.contourf(Trims, Heels, emptyArray)
plt.colorbar(label="vertical distance from CB to CG")
plt.title("Potential Energy for Geometry File "+str(gfname))
plt.xlabel("Trim (deg)")
plt.ylabel("Heel (deg)")
plt.savefig("NRG_CNTR.png")
plt.show()