from __future__ import division
from __future__ import print_function
from __future__ import with_statement

from fortranfile import FortranRead
import physconst
import os.path

from byte2human import byte2human
from time2human import time2human
from version2human import version2human

import datetime
import numpy as np

class StrData(object):
    """Interface to load KEPLER structure files."""

    def __init__(self, filename, zerotime = True):
        """
        Constructor; wants file name.
        """

        ftag = " [{}] ".format(self.__class__.__name__)
        filename = os.path.expanduser(filename)
        self.filename = os.path.expandvars(filename)
        self.file = FortranRead(self.filename)
        self._load()
        self.file.close()
        

    def _load(self, zerotime = True):
        """
        Open file, call load data, time the load, print out diagnostics.
        """
        start_time = datetime.datetime.now()
        print(self.ftag + 'Loading {} ({})'\
              .format(self.file.filename,\
                      byte2human(self.file.filesize)))
        self.data=[]
        while not self.file.eof():
            record=StrRecord(self.file)
            self.data.append(record)
        self.nmodels = len(self.data)
        end_time = datetime.datetime.now()
        load_time = end_time - start_time
        print(self.ftag + 'version {:s}'.format(version2human(self.data[0].nvers)))
        print(self.ftag + 'first model red {:>9d}'.format(int(self.data[ 0].ncyc)))
        print(self.ftag + ' last model red {:>9d}'.format(int(self.data[-1].ncyc)))
        print(self.ftag + 'data loaded in ' + time2human(load_time.total_seconds()))

        if zerotime:
            self.remove_zerotime()

    def remove_zerotime(self, verbose = True):
        """
        Detect and remove resets of time.
        """
        zerotime = np.float64(0)
        time0 = self.data[0].time
        for i in xrange(1,self.nmodels):
            if self.data[i].time < time0:
                zerotime = self.data[i-1].time
                if verbose:
                    print(self.ftag + '@ model = {:8d} zerotime was set to {:12.5g}.'\
                          .format(int(self.data[i].ncyc), \
                                  float(zerotime)))
                time0 = self.data[i].time
                self.data[i].time = time0 + zerotime


class StrRecord(object):
    """
    Individual cycle recrds of structure.

    Variables:
    ----------
    nvers - version number
    ncyc - cycle number
    timesec - current problem time (not including time offset) [s]
    dt - current time step for this cycle [s]
    jm - number of zones
    zm - mass coordinate on zone boundaries [g]
    rn - radius coordinate of zone boundaries [cm]
    tn - zone average temperature [K]
    dn - zone average density [g]
    """
    def __init__(self,file,data=True):
        self.ftag = " [{}] ".format(self.__class__.__name__)
        self.file = file
        self.load(data)
        
    def load(self,data=True):
        f = self.file
        f.load()
        self.nvers = f.get_i4()
        self.ncyc = f.get_i4()
        self.time = f.get_f8n()
        if self.nvers >= 10002:
            self.dt = f.get_f8n()
        self.jm = f.get_i4()        
        self.zm = f.get_f8an([self.jm + 1])
        self.rn = f.get_f8an([self.jm + 1])
        self.tn = f.get_f8an1d0(self.jm)
        if self.nvers >= 10001:
            self.dn = f.get_f8an1d0(self.jm)
        f.assert_eor()