from pyat.env import *
import numpy as np

def write_env_RAM(TitleEnv, freq, pos, bathym, ssp, envfil = "ram.in"):
    # envfil = "ram.in"
    Rmax = pos.rmax

    f = open(envfil, "wt")

    ndr = 1
    ndz = 1

    zmplt = pos.zmplt

    c0 = np.mean(ssp[0].cw.v)

    nPade = 8       # The number of Pade terms
    nStarter = 4    # The number of starter terms
    rs = 0.0        # The stability constraints range

    f.write("{:s}".format(TitleEnv)  + " ! Title\r\n")
    f.write("{:8.2f} ".format(freq) + \
            "{:8.2f} ".format(pos.zs) + \
            "{:8.2f} ".format(pos.zr) + \
            "\t \t \t ! Frequency (Hz) zs zr\r\n")
    f.write("{:8.2f} ".format(Rmax) + \
            "{:8.2f} ".format(pos.dr) + \
            "{:3d} ".format(ndr) + \
            "\t \t \t ! rmax dr ndr \r\n")
    f.write("{:8.2f} ".format(pos.zmax) + \
            "{:8.2f} ".format(pos.dz) + \
            "{:3d} ".format(ndz) + \
            "{:8.2f} ".format(zmplt) + \
            "\t \t \t ! zmax dz ndz zmplt\r\n")
    f.write("{:8.2f} ".format(c0) + \
            "{:1d} ".format(nPade) + \
            "{:1d} ".format(nStarter) + \
            "{:8.2f} ".format(rs) + \
            "\t \t \t ! c0 np ns rs\r\n")

    # Bathymetry
    for it in range(len(bathym.r)):
        f.write("{:8.2f} ".format(bathym.r[it]) + "{:8.2f} ".format(bathym.d[it]) + \
                " \t \t \t ! rb zb\r\n")

    # Sound speed profile
    f.write("-1 -1\r\n")
    for it in range(len(ssp)):
        # rp
        if it != 0:
            f.write("{:8.2f} \t \t \t ! rp\r\n".format(ssp[it].rp))
        # cw
        for ii in range(len(ssp[it].cw.z)):
            f.write("{:8.2f} ".format(ssp[it].cw.z[ii]) + "{:8.2f} ".format(ssp[it].cw.v[ii]) + \
                    " \t \t \t ! z cw\r\n")
        f.write("-1 -1\r\n")

        # cb
        for ii in range(len(ssp[it].cb.z)):
            f.write("{:8.2f} ".format(ssp[it].cb.z[ii]) + "{:8.2f} ".format(ssp[it].cb.v[ii]) + \
                    " \t \t \t ! z cb\r\n")
        f.write("-1 -1\r\n")

        # rhob
        for ii in range(len(ssp[it].rhob.z)):
            f.write("{:8.2f} ".format(ssp[it].rhob.z[ii]) + "{:8.2f} ".format(ssp[it].rhob.v[ii]) + \
                    " \t \t \t ! z rhob\r\n")
        f.write("-1 -1\r\n")

        # attn
        for ii in range(len(ssp[it].attn.z)):
            f.write("{:8.2f} ".format(ssp[it].attn.z[ii]) + "{:8.2f} ".format(ssp[it].attn.v[ii]) + \
                    " \t \t \t ! z attn\r\n")
        f.write("-1 -1\r\n")


def readTL_RAM(filePath=None):
    if filePath != None:
        f = open(filePath + "\\tl.grid")
    else:
        f = open("tl.grid")

    rec_marker1 = np.fromfile(f, dtype=np.int32, count=1)
    if rec_marker1 == 60:
        datatype = np.int32
    else:
        rec_maker1 = np.fromfile(f, dtype=np.int32, count=1)
        datatype = np.int64


    freq = np.fromfile(f, dtype=np.float32, count=1)
    zs = np.fromfile(f, dtype=np.float32, count=1)
    zr = np.fromfile(f, dtype=np.float32, count=1)
    rmax = np.fromfile(f, dtype=np.float32, count=1)
    dr = np.fromfile(f, dtype=np.float32, count=1)
    ndr = np.fromfile(f, dtype=np.int32, count=1)
    zmax = np.fromfile(f, dtype=np.float32, count=1)
    dz = np.fromfile(f, dtype=np.float32, count=1)
    ndz = np.fromfile(f, dtype=np.int32, count=1)
    zmplt = np.fromfile(f, dtype=np.float32, count=1)
    c0 = np.fromfile(f, dtype=np.float32, count=1)
    nPade = np.fromfile(f, dtype=np.int32, count=1)
    nStarter = np.fromfile(f, dtype=np.int32, count=1)
    rs = np.fromfile(f, dtype=np.float32, count=1)
    lz = np.fromfile(f, dtype=np.int32, count=1)

    rec_marker2 = np.fromfile(f, dtype=datatype, count=1)

    rr = np.arange(ndr*dr, rmax+ndr*dr, ndr*dr)
    rd = ndz * dz * (np.arange(lz) + 1)
    Nrr = len(rr)
    Nrd = len(rd)

    tl = np.zeros([Nrr, Nrd],dtype=np.float32)
    for it in range(Nrr):
        marker1 = np.fromfile(f, dtype=datatype, count=1)
        tl[it][:] = np.fromfile(f, dtype=np.float32, count=Nrd)
        marker2 = np.fromfile(f, dtype=datatype, count=1)

    tl = np.transpose(tl)

    return [rd, rr, tl]


class PosRAM:
    def __init__(self, zs, zr, rmax, zmax, dr, dz, zmplt):
        self.zs = zs
        self.zr = zr
        self.rmax = rmax
        self.zmax = zmax
        self.dr = dr
        self.dz = dz
        self.zmplt = zmplt


class SSPRAM:
    def __init__(self, rp, cw_p, cb_p, rhob_p, attn_p):
        self.rp = rp
        self.cw = cw_p
        self.cb = cb_p
        self.rhob = rhob_p
        self.attn = attn_p


class PROF:
    def __init__(self, depth, value):
        self.z = depth
        self.v = value