from typing import List, Dict, Any, Tuple
from pathlib import Path
import re
import numpy as np
import pandas as pd
from xml.etree.ElementTree import Element  # noqa: S405
import defusedxml.ElementTree as ET
from resistics.time import ChanMetadata
from resistics.calibrate import SensorCalibrationReader, CalibrationData
from resistics.spectra import SpectraMetadata

class SensorCalibration_RSP_RSPX_Base(SensorCalibrationReader):
    """Base class for RSP and RSPX calibration data readers"""

    file_str: str = "Metronix_Coil-----TYPE-$sensor_$chopper-ID-$serial$extension"
    """The file string to search for. Various parameters will be substituted"""

    def _get_path(self, dir_path: Path, metadata: SpectraMetadata, chan: str) -> Path:
        """
        Get the path to the calibration file

        Parameters
        ----------
        dir_path : Path
            The directory  path to look for calibration files
        metadata : SpectraMetadata
            SpectraMetadata with data information
        chan : str
            The channel to calibrate

        Returns
        -------
        Path
            The path to the calibration file
        """
        chan_metadata = metadata.chans_metadata[chan]
        chopper_str = "LF" if chan_metadata.chopper else "HF"
        sensor_str = re.sub("[^0-9]", "", chan_metadata.sensor)
        sensor_str = f"{int(sensor_str):03d}"
        serial_str = f"{int(chan_metadata.serial):06d}"
        file_name = self.file_str.replace("$sensor", sensor_str)
        file_name = file_name.replace("$serial", serial_str)
        file_name = file_name.replace("$chopper", chopper_str)
        file_name = file_name.replace("$extension", self.extension)
        return dir_path / file_name

    def _get_chopper(self, file_path: Path) -> bool:
        """Get whether the calibration is chopper on or off"""
        if "LF" in file_path.stem or "BB" in file_path.stem:
            return True
        return False

class SensorCalibrationRSP(SensorCalibration_RSP_RSPX_Base):
    """
    Reader for RSP calibration files

    RSP data is in units:

    - F [Hz]
    - Magnitude [mv/nT]
    - Phase [deg]

    Data is returned with units:

    - F [Hz]
    - Magnitude [mV/nT]
    - Phase [radians]

    The static gain for RSP files is applied to the magnitude as it is read in
    """

    extension: str = ".RSP"

    def read_calibration_data(
        self, file_path: Path, chan_metadata: ChanMetadata
    ) -> CalibrationData:
        """
        Read data from a RSP calibration file

        Parameters
        ----------
        file_path : Path
            The file path of the calibration file
        chan_metadata : ChanMetadata
            The channel metadata for the channel to be calibrated

        Returns
        -------
        CalibrationData
            The calibration data
        """
        with file_path.open("r") as f:
            lines = f.readlines()
        lines = [x.strip() for x in lines]
        data_dict = self._read_metadata(lines)
        data_dict["chopper"] = self._get_chopper(file_path)
        df = self._read_data(lines)
        df["magnitude"] = df["magnitude"] * data_dict["static_gain"]
        df["phase"] = df["phase"] * (np.pi / 180)
        data_dict["frequency"] = df.index.values.tolist()
        data_dict["magnitude"] = df["magnitude"].values.tolist()
        data_dict["phase"] = df["phase"].values.tolist()
        data_dict["file_path"] = file_path
        return CalibrationData(**data_dict)

    def _read_metadata(self, lines: List[str]) -> Dict[str, Any]:
        """Read the calibration file metadata"""
        serial, sensor = self._get_sensor_details(lines)
        static_gain = self._get_static_gain(lines)
        return {
            "serial": serial,
            "sensor": sensor,
            "static_gain": static_gain,
            "magnitude_unit": "mV/nT",
            "phase_unit": "radians",
        }

    def _get_sensor_details(self, lines: List[str]) -> Tuple[int, str]:
        """Get sensor details from the file"""
        serial: int = 1
        sensor: str = ""
        for line in lines:
            if "induction coil no" in line:
                split1 = line.split(":")[1]
                serial = int(split1.split("-")[0].strip())
            if "SensorType" in line:
                sensor = line.split()[1]
        return serial, sensor

    def _get_static_gain(self, lines: List[str]) -> float:
        static_gain: float = 1.0
        for line in lines:
            if "StaticGain" in line:
                static_gain = float(line.split()[1])
                return static_gain
        return static_gain

    def _read_data(self, lines: List[str]) -> pd.DataFrame:
        """Read data from calibration file"""
        read_from = self._get_read_from(lines)
        data_lines = self._get_data_lines(lines, read_from)
        data = np.array([x.split() for x in data_lines], dtype=np.float32)
        df = pd.DataFrame(data=data, columns=["frequency", "magnitude", "phase"])
        return df.set_index("frequency").sort_index()

    def _get_read_from(self, lines: List[str]) -> int:
        """Get the line number to read from"""
        for idx, line in enumerate(lines):
            if "FREQUENCY" in line:
                return idx + 2
        raise ValueError("Unable to determine location of data in file")

    def _get_data_lines(self, lines: List[str], idx: int) -> List[str]:
        """Get the data lines out of the file"""
        data_lines: List[str] = []
        while idx < len(lines) and lines[idx] != "":
            data_lines.append(lines[idx])
            idx += 1
        return data_lines

class SensorCalibrationRSPX(SensorCalibration_RSP_RSPX_Base):
    """
    Read data from RSPX calibration file

    RSPX data is in units:

    - F [Hz]
    - Magnitude [mv/nT]
    - Phase [deg]

    Data is returned with units:

    - F [Hz]
    - Magnitude [mV/nT]
    - Phase [radians]

    Static gain is applied to the magnitude
    """

    extension: str = ".RSPX"

    def read_calibration_data(
        self, file_path: Path, chan_metadata: ChanMetadata
    ) -> CalibrationData:
        """
        Read RSPX file

        Parameters
        ----------
        file_path : Path
            The file path of the calibration file
        chan_metadata : ChanMetadata
            The channel metadata for the channel to be calibrated

        Returns
        -------
        CalibrationData
            The calibration data
        """
        root = ET.parse(file_path).getroot()
        data_dict = self._read_metadata(root)
        data_dict["chopper"] = self._get_chopper(file_path)
        df = self._read_data(root)
        df["magnitude"] = df["magnitude"] * data_dict["static_gain"]
        df["phase"] = df["phase"] * (np.pi / 180)
        data_dict["frequency"] = df.index.values.tolist()
        data_dict["magnitude"] = df["magnitude"].values.tolist()
        data_dict["phase"] = df["phase"].values.tolist()
        data_dict["file_path"] = file_path
        return CalibrationData(**data_dict)

    def _read_metadata(self, root: Element) -> Dict[str, Any]:
        """Read the calibration file metadata"""
        serial, sensor = self._get_sensor_details(root)
        static_gain = self._get_static_gain(root)
        return {
            "serial": serial,
            "sensor": sensor,
            "static_gain": static_gain,
            "magnitude_unit": "mV/nT",
            "phase_unit": "radians",
        }

    def _get_sensor_details(self, root: Element) -> Tuple[int, str]:
        """Get sensor details"""
        serial: int = 1
        if root.find("SensorId") is not None:
            serial = int(root.find("SensorId").text)
        sensor: str = ""
        if root.find("SensorSpecification") is not None:
            sensor = root.find("SensorSpecification").text
        return serial, sensor

    def _get_static_gain(self, root) -> float:
        """Get the static gain"""
        static_gain: float = 1.0
        if root.find("StaticGain") is not None:
            static_gain = float(root.find("StaticGain").text)
        return static_gain

    def _read_data(self, root: Element) -> pd.DataFrame:
        """Get data in a DataFrame"""
        data = []
        for resp in root.findall("ResponseData"):
            data.append(
                [
                    np.float32(resp.get("Frequency")),
                    np.float32(resp.get("Magnitude")),
                    np.float32(resp.get("Phase")),
                ]
            )
        df = pd.DataFrame(data=data, columns=["frequency", "magnitude", "phase"])
        return df.set_index("frequency").sort_index()


import numpy as np
import pdb
from scipy.interpolate import interp1d
from scipy.stats import pearsonr
from scipy.stats import norm
import scipy.signal as signal
from matplotlib import gridspec
import matplotlib
matplotlib.use('Qt4Agg')
import matplotlib.pyplot as plt
from matplotlib.widgets import RadioButtons, Button, CheckButtons
import sys
from types import *

class Estimateline:
    '''Class to manually estimate where lines are located'''
    def __init__(self,pspec,ax5,uline):
        print 'If redshift calibration appears correct, hit "Accept and Close". '\
        'Otherwise, "right click" approx. where the '+uline+' line is in the plotted spectrum. '\
        'The program will re-correlate based on this guess.'
        self.ax5 = ax5
        self.cid3 = pspec.figure.canvas.mpl_connect('button_press_event',self.onclick)

    def on_key_press(self,event):
        if event.key == 'shift':
            self.shift_is_held = True

    def on_key_release(self, event):
        if event.key == 'shift':
            self.shift_is_held = False

    def onclick(self,event):
        if event.inaxes == self.ax5:
            if event.button == 3:
                print 'xdata=%f, ydata%f'%(event.xdata, event.ydata)
                self.lam = event.xdata
                plt.close()
            '''
            if event.button == 1:
                #if self.shift_is_held:
                #    print 'xdata=%f, ydata%f'%(event.xdata, event.ydata)
                #    self.lam = event.xdata
                #    plt.close()
                #else:
                plt.close()
            '''
        else: return

class z_est:
    def __init__(self,lower_w=3900.0,upper_w=5500.0,lower_z=0.01,upper_z=0.35,z_res=3.0e-5,skip_initial_priors=False):
        '''
        Initialize redshift estimate parameters
        '''

        #preconditions
        assert lower_w and upper_w, "wavelength bounds must have values"
        assert lower_z and upper_z, "redshift bounds must have values"
        assert (type(lower_w) == IntType or type(lower_w) == FloatType) and (type(upper_w) == IntType or \
                    type(upper_w) == FloatType), "wavelength bounds must be integers or floats"
        assert (type(lower_z) == IntType or type(lower_z) == FloatType) and (type(upper_z) == IntType or \
                    type(upper_z) == FloatType), "wavelength bounds must be integers or floats"
        assert lower_w < upper_w, "lower_w must be < upper_w"
        assert lower_z < upper_z, "lower_z must be < upper_z"

        #set class attributes
        self.lower_w = lower_w
        self.upper_w = upper_w
        self.lower_z = lower_z
        self.upper_z = upper_z
        self.z_res = z_res

        #create redshift array and initialize correlation value array
        self.ztest = np.arange(self.lower_z,self.upper_z,self.z_res)
        self.corr_val_i = np.zeros(self.ztest.size)

        #set redshift prior flag
        if skip_initial_priors:
            self.est_pre_z = '3'
            self.uline_n = 'HK'
            self.uline = 3950.0
            self.z_prior_width = 0.06
        else:
            self.est_pre_z = raw_input('(1) Use a known prior [Examples: median of known redshifts. Galaxy photoz measurements] \n'\
                                    '(2) View spectrum and specify a redshift prior \n'\
                                    '(3) No prior\n')

        #catch and correct false entry
        _est_enter = False
        self.uline_n = raw_input('What is the name of a spectral line you wish to use to identify redshift priors? '\
                                            '[Default: HK]: ')
        if not self.uline_n:
            self.uline_n = 'HK'
        self.uline = raw_input('Please list the approx. rest wavelength (in angstroms) of that line you seek to identify in your spectra '\
                                            '[Default: HK lines are at about 3950]: ')
        if self.uline:
            self.uline = np.float(self.uline)
        else:
            self.uline = 3950.0
        while not _est_enter:
            if self.est_pre_z == '1':
                self.z_prior_width = 0.06
                print 'redshift prior width has been set to',self.z_prior_width
                _est_enter = True
            elif self.est_pre_z == '2':
                self.z_prior_width = 0.06
                print 'redshift prior width has been set to',self.z_prior_width
                _est_enter = True
            elif self.est_pre_z == '3':
                self.z_prior_width = 0.06
                _est_enter = True
            else:
                self.est_pre_z = raw_input('Incorrect entry: Please enter either (1), (2), or (3).')

        #remind user to set the correct values in next step
        if self.est_pre_z == '1':
            print 'Make sure to set the gal_prior argument to the value of the known redshift prior: '\
                '[Example: z_est.redshift_estimate(gal_prior=0.1)]'

        #postconditions
        assert self.est_pre_z, "Must define redshift prior flag"
        assert self.est_pre_z == '1' or self.est_pre_z == '2' or self.est_pre_z == '3', \
                         "Incorrect string value for prior"

    def redshift_estimate(self,early_type_wave,early_type_flux,wave,Flux_science,plotlines=None,template_id='Galaxy',gal_prior=None):
        '''
        estimate redshift for object
        '''
        #manage redshift prior
        self.gal_prior = gal_prior
        self.template_number = template_id
        self.plotlines = plotlines

        #continuum subtract
        Flux_sc = Flux_science - signal.medfilt(Flux_science,171)
        early_type_flux_sc = early_type_flux - signal.medfilt(early_type_flux,171)

        #handle single redshift prior flag
        if self.est_pre_z == '1':
            if self.gal_prior:
                self.pre_z_est = self.gal_prior
            else:
                nospec = raw_input('You said you are either using a spectroscopic or photometric redshift prior. '\
                                        'You need to specify a prior value! Either enter a number in now or type (q) to exit')
                if nospec == 'q':
                    sys.exit()
                elif not nospec:
                    sys.exit()
                else:
                    self.gal_prior = np.float(nospec)
                    self.pre_z_est = self.gal_prior

        #handle user prior flag
        if self.est_pre_z == '2':
            print 'Take a look at the plotted galaxy spectrum and note, approximately, at what wavelength do you see the '+self.uline_n+' line. '\
                    'Then close the plot and enter that wavelength in angstroms.'
            plt.plot(wave,Flux_science)
            plt.xlim(self.lower_w,self.upper_w)
            plt.show()
            line_init = raw_input(self.uline_n+' approx. wavelength (A): ')
            self.pre_z_est = np.float(line_init)/self.uline - 1

        #handle no prior flag
        if self.est_pre_z == '3':
            self.pre_z_est = None

        redshift_est,cor,ztest,corr_val = self._cross_cor(self.pre_z_est,self.z_prior_width,early_type_wave,early_type_flux_sc,wave,Flux_sc)

        self.qualityval = 1
        self.first_pass = True
        self.skip_spec_flag = False
        self._GUI_display(redshift_est,ztest,corr_val,wave,Flux_science)
        #self.line_est = Estimateline(self.pspec,ax)
        #plt.show()
        try:
            self.pre_lam_est = self.line_est.lam
            self.pre_z_est = self.pre_lam_est/3950.0 - 1.0
            self.first_pass = False
            redshift_est,cor,ztest,corr_val = self._cross_cor(self.pre_z_est,self.z_prior_width,early_type_wave,early_type_flux_sc,wave,Flux_sc)
            print 'redshift est:',redshift_est
            self._GUI_display(redshift_est,ztest,corr_val,wave,Flux_science)
            redshift_est = self.spectra2.finalz
        except AttributeError:
            pass
        print 'zpy redshift estimate',redshift_est
        return redshift_est,cor,ztest,corr_val,self.qualityval

    def _cross_cor(self,z_est,unc,early_type_wave,early_type_flux,wave,Flux_sc):
        '''
        This function cross-correlates a continuum subtracted template spectrum with a continuum subtracted observed spectrum.
        It then returns an estimate of the redshift, the correlation value at that redshift, the array of redshifts tested,
        and the unnormalized correlation value.
        '''

        #loop over each possible redshift to compute correlation values
        for i in range(self.ztest.size):
            z = self.ztest[i]
            #redshift the template wavelengths
            wshift = early_type_wave*(1+z)
            #identify the wavelength diff between the lower wave limit and the redshifted template spectrum
            wavediff = np.min(wshift - self.lower_w)

            #if the limit is above the minimum wavelength of the redshifted template spectrum...
            if wavediff < 0:
                wave_range = wave[np.where((wave<self.upper_w)&(wave>self.lower_w))]
                Flux_range = Flux_sc[np.where((wave<self.upper_w)&(wave>self.lower_w))]
            #if the limit is below the minimum wavelength of the redshifted template spectrum...
            else:
                wave_range = wave[np.where((wave<self.upper_w+wavediff)&(wave>self.lower_w+wavediff))]
                Flux_range = Flux_sc[np.where((wave<self.upper_w+wavediff)&(wave>self.lower_w+wavediff))]

            #interpolate the redshifted template spectrum and estimate the flux at the observed spectrum wavelengths
            inter = interp1d(wshift,early_type_flux)
            et_flux_range = inter(wave_range)

            #calculate the pearson r correlation value between the observed and template flux
            self.corr_val_i[i] = pearsonr(et_flux_range,Flux_range)[0]

        #normalize the correlation values as a function of redshift
        corr_val = (self.corr_val_i[np.isfinite(self.corr_val_i)]+1)/np.trapz((self.corr_val_i[np.isfinite(self.corr_val_i)]+1),self.ztest[np.isfinite(self.corr_val_i)])
        self.ztest = self.ztest[np.isfinite(self.corr_val_i)]

        #multiply in prior to likelihood if specified
        self.corr_prior = np.zeros(self.ztest.size)
        if z_est:
            rv = norm(z_est,unc)
            corr_val = corr_val * rv.pdf(self.ztest)
            self.corr_prior = rv.pdf(self.ztest)

        #make redshift estimate
        redshift_est = (self.ztest[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))])[np.where(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))] == np.max(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))]))]

        #save correlation value at maximum redshift likelihood
        cor = (self.corr_val_i[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))])[np.where(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))] == np.max(corr_val[np.where((self.ztest>self.lower_z)&(self.ztest<self.upper_z))]))]

        return redshift_est[0], cor, self.ztest,corr_val

    def _GUI_display(self,redshift_est,ztest,corr_val,wave,flux_sc):
        '''Display the spectrum and reference lines.'''
        self.fig = plt.figure(figsize=(10, 8))
        gs = gridspec.GridSpec(2, 1, height_ratios=[2, 1])
        ax2 = plt.subplot(gs[0])
        ax = plt.subplot(gs[1])

        #maximize window
        figManager = plt.get_current_fig_manager()
        figManager.window.showMaximized()

        plt.subplots_adjust(top=0.96,bottom=0.1,left=0.04,right=0.80)

        ax.plot(ztest,corr_val,'b')
        pspec_corr = ax.axvline(redshift_est,color='k',ls='--')
        ax.fill_between(self.ztest,np.zeros(self.ztest.size),self.corr_prior,facecolor='grey',alpha=0.6)
        ax.set_xlabel('Redshift')
        ax.set_ylabel('Correlation')

        self.pspec, = ax2.plot(wave,flux_sc)
        ax2.set_ylim(np.min(flux_sc),np.max(flux_sc))
        ax2.set_xlim(wave[0],wave[-1])

        #plot lines
        vlin_pspecs = []
        if self.plotlines == None:
            self.plotlines = {'HKlines':[3968.5,3933.7],'emission_lines':[3725.0,4959.0,5007.0],'absorption_lines':[4102.9,4304.0,4862.0,5175.0],'sky_lines':[]}
        else: pass
        try:
            for vlin in self.plotlines['HKlines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='red'))
        except KeyError: pass
        try:
            for vlin in self.plotlines['emission_lines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='blue'))
        except KeyError: pass
        try:
            for vlin in self.plotlines['absorption_lines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='orange'))
        except KeyError: pass
        try:
            for vlin in self.plotlines['sky_lines']:
                vlin_pspecs.append(ax2.axvline(vlin*(1+redshift_est),ls='--',alpha=0.7,c='grey'))
        except KeyError: pass

        if self.first_pass:
            self.line_est = Estimateline(self.pspec,ax2,self.uline_n)

        self.fig.text(0.83, 0.9, 'Template {0}'.format(self.template_number), bbox=dict(facecolor='white', alpha=1.),fontsize=18)

        rax = plt.axes([0.83, 0.43, 0.15, 0.22])
        if self.qualityval == 1:
            radio = RadioButtons(rax, ('1 - No Clue      ','2 - Slight\n    Chance', '3 - Maybe', '4 - Probably', '5 - Clear'))
        else:
            radio = RadioButtons(rax, ('1 - No Clue      ','2 - Slight\n    Chance', '3 - Maybe', '4 - Probably', '5 - Clear'),active=1)
        def qualfunc(label):
            if label == '5 - Clear':
                self.qualityval = 5
            elif label == '4 - Probably':
                self.qualityval = 4
            elif label == '3 - Maybe':
                self.qualityval = 3
            elif label == '2 - Slight\n    Chance':
                self.qualityval = 2
            else:
                self.qualityval = 1

        radio.on_clicked(qualfunc)
        closeax = plt.axes([0.83, 0.3, 0.15, 0.1])
        button = Button(closeax, 'Accept & Close', hovercolor='0.975')
        def closeplot(event):
            plt.close()
        button.on_clicked(closeplot)
        skip_spec_ax = plt.axes([0.83, 0.94, 0.15, 0.04])
        skip_button = Button(skip_spec_ax, 'skip spectra', hovercolor='0.975')
        def skip_spec(event):
            plt.close()
            self.qualityval = 0
            self.skip_spec_flag = True
        skip_button.on_clicked(skip_spec)
        ax2.set_xlim(self.lower_w,self.upper_w)
        ax2.set_xlabel('Wavelength (A)')
        ax2.set_ylabel('Counts')

        plt.draw()

        if not self.first_pass:
            self.spectra2 = DragSpectra(vlin_pspecs,pspec_corr,redshift_est,ax2)
            self.fig.canvas.mpl_connect('motion_notify_event',self.spectra2.on_motion)
            self.fig.canvas.mpl_connect('button_press_event',self.spectra2.on_press)
            self.fig.canvas.mpl_connect('button_release_event',self.spectra2.on_release)
        plt.show()

class DragSpectra:
    '''Class to drage the spectra back and forth to match lines of interest'''
    def __init__(self,vlin_spectra,corr_spec,redshift_estimate,ax5):
        self.ax5 = ax5
        self.corr_spec = corr_spec
        self.yzs = self.corr_spec.get_data()[1]
        print 'begin shift'
        self.vlin_spectra = vlin_spectra
        self.vline_ys = vlin_spectra[0].get_data()[1]
        self.pressed = False
        self.finalz = redshift_estimate
        #figure.canvas.mpl_connect('motion_notify_event',self.on_motion)
        #figure.canvas.mpl_connect('button_press_event',self.on_press)
        #figure.canvas.mpl_connect('button_release_event',self.on_release)

    def on_motion(self,evt):
        if self.pressed:
            #dx = evt.xdata - self.mouse_x
            #print "%d %d" % (evt.xdata,self.mouse_x)
            newz = ((evt.xdata/self.mouse_x)*(1.+self.z_on_press))-1.  #((1. + (dx/self.mouse_x))*(1.+self.z0))-1.
            newxs = self.vline_lams*(evt.xdata/self.mouse_x) # equivalent to spec_x*((1+newz)/(1+z0))
            for i in np.arange(len(self.vlin_spectra)):
                self.vlin_spectra[i].set_data([newxs[i], newxs[i]], self.vline_ys)

            self.corr_spec.set_data([newz, newz], self.yzs)
            plt.draw()

    def on_press(self,evt):
        if evt.inaxes == self.ax5:
            self.mouse_x = evt.xdata
            self.z_on_press = self.corr_spec.get_data()[0][0]
            self.vline_lams = np.array([self.vlin_spectra[x].get_data()[0][0] for x in np.arange(len(self.vlin_spectra))])
            self.pressed = True
        else: return

    def on_release(self,evt):
        if evt.inaxes == self.ax5:
            self.pressed = False
            try:
                self.finalz = self.corr_spec.get_data()[0][0]
            except AttributeError:
                self.finalz = self.finalz
        else: return

if __name__ == '__main__':
    R = z_est()


# -*- coding: utf-8 -*-
"""
Created on Wed Jan 11 11:30:11 2017

@author: PRABLANC P.
"""

import numpy as np
import scipy.signal as sp

class ProsodicModificationRealTime(object):
    """Modification of the pitch of the voice with or without preservation of the spectrale envelope

    Notes: resampling step is not handled properly leading to "clic" sounds. The scipy.resample()
        function works with a Fourier-domain method which makes it impossible to deal with boundaries
        issues. It would be necessary to incorporate in the frame to resample a few samples before and after
        in order to deal with boundarie issues. A time-domain-based method would be probably more appropriate
        such as scipy.resample_poly()
    """

    def __init__(self, fs=16000.0, frame_format='Float32'):
        self.fs = np.float(fs)
        self._pitch_rate = 1.0
        self.MIN_INTENSITY_VALUE = -40
        self.fs_norm = self.fs / 16000.0
        self._mid_buffer_size = int(256.0 * self.fs_norm)
        self._buffer = 0.001*np.random.rand(5 * self._mid_buffer_size, 1)
        self.win_length = int(128.0 * self.fs_norm)  # based on fs = 16 kHz
        win = np.hanning(self.win_length*2)
        self.win_G = np.reshape(win[0:self.win_length], [self.win_length, 1])
        self.win_D = np.reshape(win[self.win_length:], [self.win_length, 1])
        self.win_s = int(256*self.fs_norm)

        # initialization related to pitch change
        self.initialize(self._pitch_rate, frame_format)

    def initialize(self, pitch_rate, frame_format='Float32'):
        #==============================================================================
        # Pitch-shifting initialization
        #==============================================================================
        self.frame_format = frame_format
        # if pitch_rate < 1, the weights can't overflow the pointer
        self.length_weight = int(120 * self.fs_norm)  # length of the weight window
        diff = np.ceil(self.win_s * pitch_rate) - self.length_weight
        if diff < 0:
            self.length_weight = int(self.length_weight + diff)
        tmp = np.arange(-self.length_weight, self.length_weight+1)
        # weight used in frame intercorrelation
        self.weight = np.ones(self.length_weight * 2 + 1)
        self.weight[0:self.length_weight] = 1 - np.abs(tmp[0:self.length_weight]) / float(self.length_weight) * 0.5
        self.weight[self.length_weight+1:] = 1 - np.abs(tmp[self.length_weight+1:]) / float(self.length_weight) * 0.5
        self.weight = self.weight.reshape([self.length_weight * 2 + 1, 1])
        self.frame_length = 1 + int(np.ceil(self.win_s * pitch_rate)) + 2*self.length_weight + self.win_s - 1
        self.shift_mod = int(np.ceil(self.win_s*pitch_rate))
        self.zeros_frame_mod = np.zeros([self.shift_mod, 1])
        if pitch_rate < 1:
            self.buffer_mod = np.zeros([5 * self.win_s, 1])
            print(pitch_rate)
        else:
            self.buffer_mod = np.zeros([5 * self.shift_mod, 1])

        #==============================================================================
        # LPC initialization
        #==============================================================================
        self.alpha = 0.97
        self.n_lpc = 16
        self.w_hamming = np.hamming(self.shift_mod)

    def pitchshifting(self, new_frame):
        new_frame = new_frame.reshape([new_frame.size, 1])
        self._buffer = self.buffershift(self._buffer,new_frame)

        frame_mod_prev = \
            self.buffer_mod[(self.shift_mod + 1 - self.length_weight) : (1 + self.shift_mod + self.length_weight + self.win_s)]

        frame_s = np.reshape(self._buffer[0:self.frame_length], [self.frame_length, 1])
        intensity = 10 * np.log10(np.sum(frame_s**2) + 1e-30)

        if intensity > self.MIN_INTENSITY_VALUE:
            ic = sp.correlate(frame_mod_prev, frame_s[0:self.win_s])
            intercorr = ic[self.win_s-2:self.win_s-1+2*self.length_weight] * self.weight
            I = intercorr.argmax()
            k_pos = I-(self.length_weight+1)
        else:
            k_pos = 0
        # Not clean, indices overflow buffer_mod size.
        # It works because overflowed values are taken into account.
        self.buffer_mod[self.shift_mod+k_pos+self.win_length:self.shift_mod+k_pos+self.frame_length] = \
            frame_s[self.win_length:self.frame_length]
        self.buffer_mod[self.shift_mod+k_pos:self.shift_mod+k_pos+self.win_length] = \
            + self.buffer_mod[self.shift_mod+k_pos:self.shift_mod+k_pos+self.win_length]*self.win_D \
            + frame_s[0:self.win_length]*self.win_G

        # shift modified buffer
        self.buffer_mod = self.buffershift(self.buffer_mod,self.zeros_frame_mod)
#        frame_residue = lpcanalysis(self.buffer_mod[:self.shift_mod])
        # resample
        frame_output = sp.resample(self.buffer_mod[:self.shift_mod], new_frame.size)
        return frame_output

    def lpcanalysis(self, frame):
        raise NotImplementedError('Not yet implemented');
#        frame_preemphasis = sp.lfilter([1, -self.alpha], 1, frame.reshape([frame.size]))
#        frame_preemphasis_windowed = frame_preemphasis[self.n_lpc:] \
#            *self.w_hamming
#        R = sp.correlate(frame_preemphasis_windowed, frame_preemphasis_windowed)  # coefficients de corrélation
#        Ri = R[window_length-1:window_length+self.n_lpc]
#
#        lpc_durbin = self.durbin(Ri)  # calcul des ai
#        ai = lpc_durbin['a']
#        ai = ai.reshape([ai.size])
#        self._lpc_coeff[:, n] = ai
#        frame_filt = sp.lfilter(ai, 1.0, frame_preemphasis)  # filtrage du signal
#        frame_residue = frame_filt[self.n_lpc:]

    def buffershift(self, buffer_loc, frame):
        buffer_loc[:-frame.size] = buffer_loc[frame.size:]
        buffer_loc[-frame.size:] = frame.reshape([frame.size, 1])
        return buffer_loc

    def str2numpy(self, frame):
        return np.fromstring(frame,self.frame_format)

    def numpy2str(self, frame):
        frame = frame.astype(self.frame_format)
        return frame.tostring()

    def durbin(self,frame_corr):
        """Perform linear predictive coding with Levinson-Durbin recursion.

        Args:
            frame_corr (numpy array): First samples of auto-correlation frame.

        Returns:
            {
                'a': a, the filter coefficient of an auto-regressive model
                'k': k, the reflexion coefficient
                'En': En, the prediction error
            }
        """
        R0 = frame_corr[0]
        frame_corr = np.reshape(frame_corr/R0, [frame_corr.size, 1])
        p = frame_corr.size - 1
        k = np.zeros([p, 1])
        a = 1
        for n in range(0, p):
            a = np.append(a, 0.)
            a = a.reshape([a.size, 1])
            r = frame_corr[0:n+2]
            En = np.sum(r*a)
            Bn = np.sum(np.flipud(r)*a)
            ki = -Bn/En
            a = a + np.flipud(a)*ki
            k[n] = ki
        En = R0*np.sum(frame_corr*a)
        return {
            'a': a,
            'k': k,
            'En': En,
            }

    def lsf2poly(lsf):
        """Convert line spectral frequencies to prediction filter coefficients
        returns a vector a containing the prediction filter coefficients from a vector lsf of line spectral frequencies.

        """
        #   Reference: A.M. Kondoz, "Digital Speech: Coding for Low Bit Rate Communications
        #   Systems" John Wiley & Sons 1994 ,Chapter 4

        # Line spectral frequencies must be real.

        lsf = np.array(lsf)

        if max(lsf) > np.pi or min(lsf) < 0:
            raise ValueError('Line spectral frequencies must be between 0 and pi.')

        p = len(lsf) # model order

        # Form zeros using the LSFs and unit amplitudes
        z  = np.exp(1.j * lsf)

        # Separate the zeros to those belonging to P and Q
        rQ = z[0::2]
        rP = z[1::2]

        # Include the conjugates as well
        rQ = np.concatenate((rQ, rQ.conjugate()))
        rP = np.concatenate((rP, rP.conjugate()))

        # Form the polynomials P and Q, note that these should be real
        Q  = np.poly(rQ);
        P  = np.poly(rP);

        # Form the sum and difference filters by including known roots at z = 1 and
        # z = -1

        if p%2:
            # Odd order: z = +1 and z = -1 are roots of the difference filter, P1(z)
            P1 = sp.convolve(P, [1, 0, -1])
            Q1 = Q
        else:
            # Even order: z = -1 is a root of the sum filter, Q1(z) and z = 1 is a
            # root of the difference filter, P1(z)
            P1 = sp.convolve(P, [1, -1])
            Q1 = sp.convolve(Q, [1,  1])

        # Prediction polynomial is formed by averaging P1 and Q1

        a = .5 * (P1+Q1)
        return a[0:-1:1] # do not return last element

    def poly2lsf(a):
        """Prediction polynomial to line spectral frequencies.

        converts the prediction polynomial specified by A,
        into the corresponding line spectral frequencies, LSF.
        normalizes the prediction polynomial by A(1).

        """

        #Line spectral frequencies are not defined for complex polynomials.

        # Normalize the polynomial

        a = np.array(a)
        if a[0] != 1:
            a/=a[0]

        if max(np.abs(np.roots(a))) >= 1.0:
            raise ValueError('The polynomial must have all roots inside of the unit circle.');

        # Form the sum and difference filters

        p  = len(a)-1   # The leading one in the polynomial is not used
        a1 = np.concatenate((a, np.array([0])))
        a2 = a1[-1::-1]
        P1 = a1 - a2        # Difference filter
        Q1 = a1 + a2        # Sum Filter

        # If order is even, remove the known root at z = 1 for P1 and z = -1 for Q1
        # If odd, remove both the roots from P1

        if p%2: # Odd order
            P, r = sp.deconvolve(P1,[1, 0 ,-1])
            Q = Q1
        else:          # Even order
            P, r = sp.deconvolve(P1, [1, -1])
            Q, r = sp.deconvolve(Q1, [1,  1])

        rP  = np.roots(P)
        rQ  = np.roots(Q)

        aP  = np.angle(rP[1::2])
        aQ  = np.angle(rQ[1::2])

        lsf = sorted(np.concatenate((-aP,-aQ)))

        return lsf

    @property
    def mid_buffer_size(self):
        return self._mid_buffer_size

    @property
    def buffer(self):
        return self._buffer

    @property
    def pitch_rate(self):
        print("pitch_rate = {}".format(self._pitch_rate))
        return self._pitch_rate

    @pitch_rate.setter
    def pitch_rate(self, new_pitch_rate):
        self._pitch_rate = new_pitch_rate
        self.initialize(self._pitch_rate, self.frame_format)


from __future__ import (absolute_import, division,
                        print_function, unicode_literals)
import numpy as np
from scipy import stats
from .profiles import GaussProfile
from .profiles import UserProfile
from .profiles import DataProfile
from ..utils.utils import make_quant, shift_t

class Pulsar(object):
    """class for pulsars

    The minimal data to instantiate a pulsar is the period, Smean, and
    pulse profile. The Profile is supplied via a :class:`PulseProfile`-like
    object.

    Parameters
    ----------

    period : float
        Pulse period (sec)

    Smean : float
        Mean pulse flux density (Jy)

    profile : :class:`PulseProfile`
        Pulse profile or 2-D pulse portrait

    name : str
        Name of pulsar
    """
    #TODO Other data could be supplied via a `.par` file.
    def __init__(self, period, Smean, profiles=None, name=None):
        self._period = make_quant(period, 's')
        self._Smean = make_quant(Smean, 'Jy')

        self._name = name

        # Assign profile class; default to GaussProfile if nothing is specified
        if profiles is None:
            self._Profiles = GaussProfile()
        else:
            self._Profiles = profiles

    def __repr__(self):
        namestr = "" if self.name is None else self.name+", "
        return "Pulsar("+namestr+"{})".format(self.period.to('ms'))

    @property
    def Profiles(self):
        return self._Profiles

    @property
    def name(self):
        return self._name

    @property
    def period(self):
        return self._period

    @property
    def Smean(self):
        return self._Smean

    def make_pulses(self, signal, tobs):
        """generate pulses from Profiles, :class:`PulsePortrait` object

        Required Args:
            signal (:class:`Signal`-like): signal object to store pulses
            tobs (float): observation time (sec)
        """
        signal._tobs = make_quant(tobs, 's')

        # init base profile at correct sample rate
        Nph = int((signal.samprate * self.period).decompose())

        # If there isn't enough frequencies in the profiles
        # And if it is a :class:`Profile` instance, reshape.
        self.Profiles.init_profiles(Nph, signal.Nchan)

        # if (signal.Nchan != self.Profiles.profiles.shape[0]
        #     and hasattr(self.Profiles, 'set_Nchan')):
        #     self.Profiles.set_Nchan(signal.Nchan)

        # select pulse generation method
        if signal.sigtype in ["RFSignal", "BasebandSignal"]:
            self._make_amp_pulses(signal)
        elif signal.sigtype is "FilterBankSignal":
            self._make_pow_pulses(signal)
        else:
            msg = "no pulse method for signal: {}".format(signal.sigtype)
            raise NotImplementedError(msg)

        # compute Smax (needed for radiometer noise level)
        #pr = self.Profiles()
        pr = self.Profiles._max_profile
        nbins = len(pr) # Think this assumes a single profile for now...
        signal._Smax = self.Smean * nbins / np.sum(pr)

    def _make_amp_pulses(self, signal):
        """generate amplitude pulses

        This method should be used for radio frequency and basebanded
        pulses.

        Parameters
        ----------

        signal : :class:`Signal`-like
            Signal object to store pulses.
        """
        # generate several pulses in time
        distr = stats.norm()

        Nsamp = int((signal.tobs * signal.samprate).decompose())
        signal.init_data(Nsamp)

        # TODO break into blocks
        # TODO phase from .par file
        # calc profile at phases
        phs = (np.arange(Nsamp) /
                (signal.samprate * self.period).decompose().value)
        phs %= 1  # clip integer part

        # convert intensity profile to amplitude!
        full_prof = np.sqrt(self.Profiles.calc_profiles(phs))

        signal._data = full_prof * distr.rvs(size=signal.data.shape)

    def _make_pow_pulses(self, signal):
        """generate a power pulse

        This method should be used for filter bank pulses

        Parameters
        ----------

        signal : :class:`Signal`-like
            Signal object to store pulses.
        """
        if signal.fold:
            # Determine how many subints to make
            if signal.sublen is None:
                signal._sublen = signal.tobs
                signal._nsub = 1
            else:
                # This should be an integer, if not, will round
                signal._nsub = int(np.round(signal.tobs / signal.sublen))

            # determine the number of data samples necessary
            signal._nsamp = int((signal.nsub*(self.period*signal.samprate)).decompose())
            # Need to make an initial empty data array
            signal.init_data(signal.nsamp)

            # Tile the profiles to number of desired subints
            sngl_prof = np.tile(self.Profiles(), signal.nsub)

            # changed to number of subints
            signal._Nfold = (signal.sublen / self.period).decompose()
            distr = stats.chi2(df=signal.Nfold)
            signal._set_draw_norm(df=signal.Nfold)

            #Why is there a second call to init_data?
            signal.init_data(sngl_prof.shape[1])
            signal._data = (sngl_prof * distr.rvs(size=signal.data.shape)
                            * signal._draw_norm)
        else:
            # fold is false and will make single pulses
            signal._sublen = self.period
            # This should be an integer, if not, will round; may not be exact
            signal._nsub = int(np.round((signal.tobs / signal.sublen).decompose()))

            # generate several pulses in time
            distr = stats.chi2(df=1)
            signal._set_draw_norm(df=1)

            signal._nsamp = int((signal.tobs * signal.samprate).decompose())
            signal.init_data(signal.nsamp)

            # TODO break into blocks
            # TODO phase from .par file
            # calc profiles at phases
            phs = (np.arange(signal.nsamp) /
                   (signal.samprate * self.period).decompose().value)
            phs %= 1  # clip integer part
            full_prof = self.Profiles.calc_profiles(phs,signal.Nchan)

            signal._data = (full_prof * distr.rvs(size=signal.data.shape)
                            * signal._draw_norm)

    def null(self, signal, null_frac, length=None, frequency=None):
        """
        Function to simulate pulsar pulse nulling. Given some nulling fraction,
        will replace simulated pulses with noise until nulling fraction is met.
        This function should only be run after running any ism or other delays
        have been added, e.g. disperison, FD, profile evolution, etc., but
        should be run before adding the radiometer noise ('telescope.observe()`),
        if nulling is desired.

        Parameters
        ----------

        signal [class] : signal class containing the simulated pulses
        null_frac [float] : desired pulsar nulling fraction, given as a
            decimal; range of 0.0 to 1.0.
        length [float] : desired length of each null in seconds. If not given,
                         will randomly null pulses. Default is None.
        frequency [float] : frequency of pulse nulling, e.g. how often the pulsar
                            nulls per hour. E.g. if frequency is 2, then the
                            pulsar will null twice per hour for some length
                            of time. If not given, will randomly null pulses.
                            Default is None.
        """
        # Determine how many pulses need to be nulled out
        null_pulses = int(np.round(signal.nsub * null_frac)) # needs to be int, may not be exact
        # Get the number of bins per pulse
        Nph = int((signal.samprate * self.period).decompose())
        # determine the off pulse window
        opw = self.Profiles._calcOffpulseWindow(Nphase = Nph)
        # define the random noise distribution
        if signal.fold:
            distr = stats.chi2(df=signal.Nfold)
        else:
            distr = stats.chi2(df=1)
        # have a test ditribution to determine null bins if Nfold is 1
        if not signal.fold or signal.Nfold < 100:
            check_distr = stats.chi2(df=100)
        else:
            check_distr = stats.chi2(df=signal.Nfold)
        # Figure out how off-center the pulses are
        shift_val = Nph//2 - np.where(signal.data[0,:Nph]==np.max(signal.data[0,:Nph]))[0]
        # Now null randomly if no length or frequency is given
        if length==None and frequency==None:
            # randomly draw pulse numbers to null
            rand_pulses = np.random.choice(signal.nsub, null_pulses, replace=False)
            # replace each pulse number with random noise
            if signal.delay == None:
                for p in rand_pulses:
                    # convert to bins
                    null_bins = np.arange(Nph*p, Nph*(p+1)) + shift_val
                    # make sure we don't go beyond the data array length
                    null_bins = null_bins[null_bins<np.shape(signal.data)[1]]
                    # replace pulses with noise
                    if len(null_bins) != Nph:
                        noise = (distr.rvs(size=len(null_bins)) * signal._draw_norm)
                    else:
                        noise = (distr.rvs(size=Nph) * signal._draw_norm)
                    # if no extra delays have been added to the signal
                    signal._data[:,null_bins] = noise * np.mean(self.Profiles._max_profile[opw.astype(int)])
            # if if signal has been delayed (e.g dispersed, etc.)
            else:
                # First initialize new array
                null_array = np.zeros(np.shape(signal.data))
                for p in rand_pulses:
                    #null_bins = np.arange(Nph*p, Nph*(p+1))
                    null_bins = np.arange(Nph*p, Nph*(p+1)) + shift_val
                    # make sure we don't go beyond the data array length
                    null_bins = null_bins[null_bins<np.shape(signal.data)[1]]
                    # replace the appropriate arrays
                    if len(null_bins) != Nph:
                        null_array[:,null_bins] = (check_distr.rvs(size=len(null_bins)) * signal._draw_norm)
                    else:
                        null_array[:,null_bins] = (check_distr.rvs(size=Nph) * signal._draw_norm)
                # now shift the data
                freq_array = signal._dat_freq
                shift_dt = (1/signal._samprate).to('ms')
                for ii, freq in enumerate(freq_array):
                    null_array[ii,:] = shift_t(null_array[ii,:],
                                                 signal.delay[ii].value,
                                                 dt=shift_dt.value)
                # Now replace pulses with noise
                off_pulse_mean = np.mean(self.Profiles._max_profile[opw.astype(int)])
                noise_shape = np.shape(np.where(null_array>1))[1]
                noise = (distr.rvs(size=noise_shape) * signal._draw_norm)
                signal._data[np.where(null_array>1)] = noise*off_pulse_mean

        else:
            raise NotImplementedError("Length and Frequency not been implimented yet")


"""
resp_optimizer.py
Implementation of open-source vrsion of RESP method

"""
import os, sys
import numpy as np
from shutil import copyfile
from respyte.readinp_esp import *
from respyte.readmol import * # read molecules and assign (what?)
from respyte.molecule import * # copied from forcebalance package and modified
from respyte.gen_grid_pts import * # take either rdkit ol or oechem mol to generate grid points around molecules (modified CIB's code)
from respyte.engine import *

def main():
    # cwd = current working directory in which input folder exists
    cwd = os.getcwd()
    # check if the current working idr contains input folder
    print('\n')
    print('---------------------------------------------------------')
    print(' 1. Reading  input files and folders.                    ')
    print('---------------------------------------------------------')
    if os.path.isdir("%s/input" % cwd):
        print(' Found the input folder. Now read input/input.yml')
    else:
        print(' Failed to find input folder. Should have input(folder) containing input.yml and molecules(folder)')
    # read input.yml and generate gridOption.
    inp = Input('%s/input/input.yml' % (cwd))
    gridOptions = inp.gridOptions

    # Visit all subfolder molecule/ moli/confj and generate grid.dat
    for idx, i in enumerate(inp.nmols):
        molN = 'mol%d' % (idx+1)
        wkd = '%s/input/molecules/%s' % (cwd, molN)
        for j in range(i):
            confN = 'conf%d' % (j+1)
            path = wkd + '/%s' % (confN)
            # gen molecule object from pdb file
            pdbfile = path + '/%s_%s.pdb' % (molN, confN)
            mol2file = path + '/%s_%s.mol2' % (molN, confN)
            xyzfile = path + '/%s_%s.xyz' % (molN, confN)
            if os.path.isfile(pdbfile):
                coordfile = pdbfile
                #mol = Molecule(pdbfile)
            elif os.path.isfile(mol2file):
                coordfile = mol2file
                #mol = Molecule(mol2file)
            elif os.path.isfile(xyzfile):
                print(" Warning: Are you using xyz files? xyz format is not recommendable.")
                print("          But no worries! It will convert your xyz files to pdb with single residue named 'MOL'.")
                xyztopdb(xyzfile)
                coordfile = pdbfile
            else:
                raise RuntimeError(" Coordinate file does not exist! ")
            # Assign radii
            radii = inp.settings['radii']
            if inp.cheminformatics == 'openeye':
                mol = ReadOEMolFromFile(coordfile)
                mol = assignRadii(mol, 'OEMol', radii)
                moltype = 'OEMol'
            elif inp.cheminformatics == 'rdkit':
                mol = ReadRdMolFromFile(coordfile)
                mol = assignRadii(mol, 'RDMol', radii)
                moltype = 'RDMol'
            else:
                mol = ReadMolFromFile(coordfile)
                mol = assignRadii(mol, 'FBMol', radii)
                moltype = 'FBMol'
            # Generate points
            print('\n')
            print('---------------------------------------------------------')
            print(' 2. Generating grid points around each molecules.        ')
            print('---------------------------------------------------------')

            if moltype is 'FBMol':
                raise NotImplementedError('Grid gen using fbmol not implemented yet! Sorry!')
            if gridOptions['gridType'] == 'MSK' or gridOptions['gridType']=='extendedMsk':
                # generate Mers-Singh-Kollman Connolly surfaces at 1.4, 1.6, 1.8 and 2.0 * vdW radii
                allPts = GenerateMSKShellPts(mol, gridOptions, moltype)
            else:
                # Generate a face-centered cubic grid shell around the molecule using gridOptions
                allPts = [GenerateGridPointSetAroundOEMol(mol, gridOptions, moltype)]
            print('Total points:', np.sum([len(i) for i in allPts]))

            # Write grid.dat
            print('\n')
            print('---------------------------------------------------------')
            print(' Psi calculations are done. Will create espf files.      ')
            print('---------------------------------------------------------')
            tmppath = '%s/tmp' % path
            if inp.settings['forcegen'] == 'N' and os.path.isfile('%s/grid.dat' % path):
                if os.path.isdir(tmppath):
                    print('Rename pre-existing tmp dir into tmp.0 and re-make tmp dir.')
                    os.rename(tmppath, '%s/tmp.0' % path)
                    os.mkdir(tmppath)
                    copyfile('%s/grid.dat' % path, '%s/grid.dat' % tmppath)
                else:
                    os.mkdir(tmppath)
                    # copy grid.dat into tmp dir
                    copyfile('%s/grid.dat' % path, '%s/grid.dat' % tmppath)
                print(' Pre-existing grid.dat has been copied into tmp for Psi4 calculation.')
            else:
                if os.path.isdir(tmppath):
                    print(' Rename pre-existing tmp dir into tmp.0 and re-make tmp dir.')
                    os.rename(tmppath, '%s/tmp.0' % path)
                os.mkdir(tmppath)
                ofs = open('%s/grid.dat' % tmppath,'w')

                for pts in allPts: #changed (for msk-based scheme, allPts is a list of lists)
                    for pt in pts:
                        ofs.write( '{0:10.6f} {1:10.6f} {2:10.6f}\n'.format(pt[0], pt[1], pt[2]) )
                ofs.close()
                print(' grid.dat is successfully generated in %s.' % tmppath)
            # Run engine to run Psi4 calculation inside tmp directory.
            engine = EnginePsi4()
            if 'method' in inp.settings:
                mthd = inp.settings['method']
            else:
                mthd = 'hf'
            if 'basis' in inp.settings:
                bss = inp.settings['basis']
            else:
                bss = '6-31g*'

            chg = inp.charges[idx]

            ### Need to add 'solvation' setting!
            if 'pcm' in inp.settings:
                if inp.settings['pcm'] is 'N':
                    solv = None
                elif inp.settings['pcm'] is 'Y':
                    solv = inp.settings['solvent']
            else:
                solv = None

            engine.write_input(coordfile, method = mthd, basis = bss, charge = chg, solvent=solv, job_path = tmppath)
            engine.espcal(job_path= tmppath)
            griddat = tmppath + '/grid.dat'
            espdat = tmppath + '/grid_esp.dat'
            efdat = tmppath + '/grid_field.dat'
            espfoutput = tmppath + '/%s_%s.espf' % (molN, confN)
            engine.genespf(griddat, espdat, efdat, espfoutput)
            # copy espf file generated in moli/confj/tmp to moli/confj/
            copyfile(espfoutput, path + '/%s_%s.espf' % (molN, confN))
            copyfile(griddat, path + '/grid.dat')
    print('\n')
    print('---------------------------------------------------------')
    print('       Done! Hope it helps, Will see you again:)         ')
    print('---------------------------------------------------------')

if __name__ == "__main__":
    main()


# encoding:utf-8
import sys

sys.path.append("..")
import numpy as np
from collections import defaultdict
from random import randint
from random import shuffle,choice
from math import log
import gensim.models.word2vec as w2v
from prettyprinter import cpprint

from mf import MF
from reader.trust import TrustGetter
from utility.matrix import SimMatrix
from utility.similarity import cosine
from utility import util

class CUNE(MF):
    """
    docstring for CUNE

    Zhang et al. Collaborative User Network Embedding for Social Recommender Systems. SDM
    """

    def __init__(self):
        super(CUNE, self).__init__()
        self.config.lambdaP = 0.01
        self.config.lambdaQ = 0.01
        self.config.alpha = 0.01
        self.config.isEarlyStopping = True
        self.tg = TrustGetter()
        self.config.walkCount = 30
        self.config.walkLength = 20
        self.config.walkDim = 20
        self.config.winSize = 5
        self.config.topK = 50

    def init_model(self, k):
        super(CUNE, self).init_model(k)
        self.user_sim = SimMatrix()
        self.generate_cu_net()
        self.deep_walk()
        self.compute_social_sim()

    def generate_cu_net(self):
        print('Building collaborative user network...')
        itemNet = {}
        for item in self.rg.trainSet_i:
            if len(self.rg.trainSet_i[item])>1:
                itemNet[item] = self.rg.trainSet_i[item]

        filteredRatings = defaultdict(list)
        for item in itemNet:
            for user in itemNet[item]:
                if itemNet[item][user] > 0:
                    filteredRatings[user].append(item)

        self.CUNet = defaultdict(list)

        for user1 in filteredRatings:
            s1 = set(filteredRatings[user1])
            for user2 in filteredRatings:
                if user1 != user2:
                    s2 = set(filteredRatings[user2])
                    weight = len(s1.intersection(s2))
                    if weight > 0:
                        self.CUNet[user1]+=[user2] # * weight

        # cpprint(self.CUNet)
        pass
    def deep_walk(self):
        print('Generating random deep walks...')
        self.walks = []
        self.visited = defaultdict(dict)
        for user in self.CUNet:
            for t in range(self.config.walkCount):
                path = [str(user)]
                lastNode = user
                for i in range(1,self.config.walkLength):
                    nextNode = choice(self.CUNet[lastNode])
                    count=0
                    while(nextNode in self.visited[lastNode]):
                        nextNode = choice(self.CUNet[lastNode])
                        #break infinite loop
                        count+=1
                        if count==self.config.walkLength: # 10
                            break
                    path.append(str(nextNode))
                    self.visited[user][nextNode] = 1
                    lastNode = nextNode
                self.walks.append(path)
        # shuffle(self.walks)
        # cpprint(self.walks)
        print('Generating user embedding...')
        self.model = w2v.Word2Vec(self.walks, size=self.config.walkDim, window=5, min_count=0, iter=3)
        print('User embedding generated.')
        pass

    def compute_social_sim(self):
        print('Constructing similarity matrix...')
        # self.W = np.zeros((self.rg.get_train_size()[0], self.config.walkDim))
        self.topKSim = defaultdict(dict)
        i = 0
        for user1 in self.CUNet:
            sims = {}
            for user2 in self.CUNet:
                if user1 != user2:
                    wu1 = self.model[str(user1)]
                    wu2 = self.model[str(user2)]
                    sims[user2]=cosine(wu1,wu2) #若为空咋整
            self.topKSim[user1] = sorted(sims.items(), key=lambda d: d[1], reverse=True)[:self.config.topK]
            i += 1
            if i % 200 == 0:
                print('progress:', i, '/', len(self.CUNet))
        # print(self.topKSim)
        #构建被关注列表
        print('Constructing desimilarity matrix...')
        self.topKSimBy = defaultdict(dict)
        for user in self.topKSim:
            users=self.topKSim[user]
            for user2 in users:
                self.topKSimBy[user2[0]][user] = user2[1]
        print('Similarity matrix finished.')

    def train_model(self, k):
        super(CUNE, self).train_model(k)
        iteration = 0
        while iteration < self.config.maxIter:
            self.loss = 0
            for index, line in enumerate(self.rg.trainSet()):
                user, item, rating = line
                u = self.rg.user[user]
                i = self.rg.item[item]
                error = rating - self.predict(user, item)
                self.loss += 0.5 * error ** 2
                p, q = self.P[u], self.Q[i]

                social_term_p, social_term_loss = np.zeros((self.config.factor)), 0.0
                followees = self.topKSim[user] #self.tg.get_followees(user) #self.topKSim[user]
                # print(followees)
                for followee in followees:
                    if self.rg.containsUser(followee[0]):
                        # s = self.user_sim[user][followee]
                        uf = self.P[self.rg.user[followee[0]]]
                        social_term_p += followee[1]* (p - uf)
                        social_term_loss += followee[1]* ((p - uf).dot(p - uf))

                social_term_m = np.zeros((self.config.factor))
                followers = self.topKSimBy[user]
                followers = sorted(followers.items(), key=lambda d: d[1], reverse=True)[:self.config.topK]
                for follower in followers:
                    if self.rg.containsUser(follower[0]):
                        ug = self.P[self.rg.user[follower[0]]]
                        social_term_m += follower[1]*(p - ug)

                # update latent vectors
                self.P[u] += self.config.lr * (
                        error * q - self.config.alpha * (social_term_p + social_term_m) - self.config.lambdaP * p)
                self.Q[i] += self.config.lr * (error * p - self.config.lambdaQ * q)

                self.loss += 0.5 * self.config.alpha * social_term_loss

            self.loss += 0.5 * self.config.lambdaP * (self.P * self.P).sum() + 0.5 * self.config.lambdaQ * (
                    self.Q * self.Q).sum()

            iteration += 1
            if self.isConverged(iteration):
                break

if __name__ == '__main__':
    # srg = CUNE()
    # srg.train_model(0)
    # coldrmse = srg.predict_model_cold_users()
    # print('cold start user rmse is :' + str(coldrmse))
    # srg.show_rmse()

    rmses = []
    maes = []
    cunemf = CUNE()
    # cunemf.init_model(0)
    # cunemf.generate_cu_net()
    # cunemf.deep_walk()
    # print(bmf.rg.trainSet_u[1])
    cunemf.config.k_fold_num = 5
    for i in range(cunemf.config.k_fold_num):
        print('the %dth cross validation training' % i)
        cunemf.train_model(i)
        rmse, mae = cunemf.predict_model()
        rmses.append(rmse)
        maes.append(mae)
    rmse_avg = sum(rmses) / cunemf.config.k_fold_num
    mae_avg = sum(maes) / cunemf.config.k_fold_num
    print("the rmses are %s" % rmses)
    print("the maes are %s" % maes)
    print("the average of rmses is %s " % rmse_avg)
    print("the average of maes is %s " % mae_avg)


from __future__ import division, absolute_import, print_function

from builtins import range

import numpy as np
import scipy.interpolate as interpolate
import scipy.integrate as integrate
import os
import sys
from pkg_resources import resource_exists
from pkg_resources import resource_filename
from pkg_resources import resource_listdir

try:
    import fitsio
    fits_package = 'fitsio'
except ImportError:
    try:
        import astropy.io.fits as pyfits
        fits_package = 'pyfits'
    except ImportError:
        try:
            import pyfits
            fits_package = 'pyfits'
        except ImportError:
            pass

from .modtranGenerator import ModtranGenerator

from .sharedNumpyMemManager import SharedNumpyMemManager as snmm
from .fgcmLogger import FgcmLogger

class FgcmAtmosphereTable(object):
    """
    Class to generate, store, and interpolate pre-computed atmospheres.

    parameters
    ----------
    atmConfig: dict
       Dictionary with FgcmAtmosphereTable config variables
    atmosphereTableFile: str, optional
       Name of the table file
    fgcmLog: FgcmLog, optional
       Logger
    """

    def __init__(self, atmConfig, atmosphereTableFile=None, fgcmLog=None):

        if fgcmLog is None:
            self.fgcmLog = FgcmLogger('dummy.log', 'INFO', printLogger=True)
        else:
            self.fgcmLog = fgcmLog

        if atmosphereTableFile is not None:
            # We have an atmosphereTableFile, so read that
            self.atmosphereTableFile = atmosphereTableFile
            self.atmConfig = atmConfig
        else:
            # get modTran stuff ready since we're going to be generating
            # also check the config
            requiredKeys = ['elevation',
                            'pmbRange','pmbSteps',
                            'pwvRange','pwvSteps',
                            'o3Range','o3Steps',
                            'tauRange','tauSteps',
                            'alphaRange','alphaSteps',
                            'zenithRange','zenithSteps',
                            'pmbStd','pwvStd','o3Std',
                            'tauStd','alphaStd','airmassStd']

            for key in requiredKeys:
                if key not in atmConfig:
                    raise ValueError("Required %s not in atmConfig" % (key))
                if 'Range' in key:
                    if len(atmConfig[key]) != 2:
                        raise ValueError("%s must have 2 elements" % (key))

            self.atmConfig = atmConfig

            self.elevation = self.atmConfig['elevation']

            self.modGen = ModtranGenerator(self.elevation)
            self.pmbElevation = self.modGen.pmbElevation

            self.pmbStd = self.atmConfig['pmbStd']
            self.pwvStd = self.atmConfig['pwvStd']
            self.lnPwvStd = np.log(self.pwvStd)
            self.o3Std = self.atmConfig['o3Std']
            self.tauStd = self.atmConfig['tauStd']
            self.lnTauStd = np.log(self.tauStd)
            self.alphaStd = self.atmConfig['alphaStd']
            self.secZenithStd = self.atmConfig['airmassStd']
            self.zenithStd = np.arccos(1./self.secZenithStd)*180./np.pi

            if ('lambdaRange' in self.atmConfig):
                self.lambdaRange = np.array(self.atmConfig['lambdaRange'])
            else:
                self.lambdaRange = np.array([3000.0,11000.0])

            if ('lambdaStep' in self.atmConfig):
                self.lambdaStep = self.atmConfig['lambdaStep']
            else:
                self.lambdaStep = 0.5

            if ('lambdaNorm' in self.atmConfig):
                self.lambdaNorm = self.atmConfig['lambdaNorm']
            else:
                self.lambdaNorm = 7750.0

        self.o2Interpolator = None
        self.o3Interpolator = None
        self.rayleighInterpolator = None
        self.lnPwvInterpolator = None

    @staticmethod
    def getAvailableTables():
        """
        Get list of installed tables

        parameters
        ----------
        None

        returns
        -------
        List of installed table names
        """

        try:
            files = resource_listdir(__name__,'data/tables/')
        except:
            raise IOError("Could not find associated data/tables path!")

        # build a dictionary and put in the key details of each...

        availableTables = {}

        for f in files:
            availableTables[os.path.basename(f)] = FgcmAtmosphereTable.getInfoDict(f)

        return availableTables

    @staticmethod
    def getInfoDict(atmosphereTableFile):
        """
        Get information dictionary on a table file

        parameters
        ----------
        atmosphereTableFile: str
           Name of the table file

        returns
        -------
        infoDict: dict
           Dictionary with key parameters describing the table
        """

        if fits_package == 'fitsio':
            parStruct = fitsio.read(atmosphereTableFile, ext='PARS')
        elif fits_package == 'pyfits':
            parStruct = pyfits.getdata(atmosphereTableFile, ext=('PARS', 1))
        else:
            raise IOError("Reading atmosphere tables not supported without fitsio/astropy.io.fits/pyfits")

        infoDict = {'elevation':parStruct['ELEVATION'][0],
                    'pmbRange':[parStruct['PMB'][0][0], parStruct['PMB'][0][-1]],
                    'pmbSteps':parStruct['PMB'][0].size,
                    'pwvRange':[np.exp(parStruct['LNPWV'][0][0]),
                                np.exp(parStruct['LNPWV'][0][-1])],
                    'pwvSteps':parStruct['LNPWV'][0].size,
                    'o3Range':[parStruct['O3'][0][0], parStruct['O3'][0][-1]],
                    'o3Steps':parStruct['O3'][0].size,
                    'tauRange':[np.exp(parStruct['LNTAU'][0][0]), np.exp(parStruct['LNTAU'][0][-1])],
                    'tauSteps':parStruct['LNTAU'][0].size,
                    'alphaRange':[parStruct['ALPHA'][0][0], parStruct['ALPHA'][0][-1]],
                    'alphaSteps':parStruct['ALPHA'][0].size,
                    'zenithRange':[np.rad2deg(np.arccos(1./parStruct['SECZENITH'][0][0])),
                                   np.rad2deg(np.arccos(1./parStruct['SECZENITH'][0][-1]))],
                    'zenithSteps':parStruct['SECZENITH'][0].size,
                    'pmbStd':parStruct['PMBSTD'][0],
                    'pwvStd':parStruct['PWVSTD'][0],
                    'o3Std':parStruct['O3STD'][0],
                    'tauStd':parStruct['TAUSTD'][0],
                    'alphaStd':parStruct['ALPHASTD'][0],
                    'airmassStd':parStruct['AIRMASSSTD'][0]}

        return infoDict

    @classmethod
    def initWithTableName(cls, atmosphereTableName):
        """
        Initialize FgcmAtmosphereTable with a table name

        parameters
        ----------
        atmosphereTableName: str
           Name of atmosphere table

        returns
        -------
        FgcmAtmosphereTable: FgcmAtmosphereTable object

        notes
        -----
        Raises IOError if atmosphereTableName couldn't be found
        """

        # first, check if we have something in the path...
        if os.path.isfile(atmosphereTableName):
            atmosphereTableFile = os.path.abspath(atmosphereTableName)
            print("Found atmosphereTableName: %s" % (atmosphereTableName))
        else:
            # allow for a name with or without .fits extension
            if resource_exists(__name__,'data/tables/%s' % (atmosphereTableName)):
                testFile = 'data/tables/%s' % (atmosphereTableName)
            elif resource_exists(__name__,'data/tables/%s.fits' % (atmosphereTableName)):
                testFile = 'data/tables/%s.fits' % (atmosphereTableName)
            else:
                raise IOError("Could not find atmosphereTableName (%s) in the path or in data/tables/" % (atmosphereTableName))
            try:
                atmosphereTableFile = resource_filename(__name__,testFile)
            except:
                raise IOError("Error finding atmosphereTableName (%s)" % (atmosphereTableName))

        # will set self.atmConfig
        print("Found atmosphere file: %s" % (atmosphereTableFile))
        atmConfig = FgcmAtmosphereTable.getInfoDict(atmosphereTableFile)

        return cls(atmConfig, atmosphereTableFile=atmosphereTableFile)

    def loadTable(self):
        """
        Load atmosphere table

        parameters
        ----------
        None

        returns
        -------
        None
        """

        # note that at this point these are the only options
        if fits_package == 'fitsio':
            parStruct = fitsio.read(self.atmosphereTableFile, ext='PARS')
        else:
            parStruct = pyfits.getdata(self.atmosphereTableFile, ext=('PARS', 1))

        self.elevation = parStruct['ELEVATION'][0]
        self.pmbElevation = parStruct['PMBELEVATION'][0]
        self.pmbStd = parStruct['PMBSTD'][0]
        self.pwvStd = parStruct['PWVSTD'][0]
        self.o3Std = parStruct['O3STD'][0]
        self.tauStd = parStruct['TAUSTD'][0]
        self.alphaStd = parStruct['ALPHASTD'][0]
        self.secZenithStd = parStruct['AIRMASSSTD'][0]
        self.lambdaRange = parStruct['LAMBDARANGE'][0]
        self.lambdaStep = parStruct['LAMBDASTEP'][0]
        self.lambdaNorm = parStruct['LAMBDANORM'][0]

        self.atmLambda = parStruct['ATMLAMBDA'][0]
        self.atmStdTrans = parStruct['ATMSTDTRANS'][0]

        self.pmb = parStruct['PMB'][0]
        self.pmbDelta = parStruct['PMBDELTA'][0]
        self.lnPwv = parStruct['LNPWV'][0]
        self.lnPwvDelta = parStruct['LNPWVDELTA'][0]
        self.o3 = parStruct['O3'][0]
        self.o3Delta = parStruct['O3DELTA'][0]
        self.lnTau = parStruct['LNTAU'][0]
        self.lnTauDelta = parStruct['LNTAUDELTA'][0]
        self.alpha = parStruct['ALPHA'][0]
        self.alphaDelta = parStruct['ALPHADELTA'][0]
        self.secZenith = parStruct['SECZENITH'][0]
        self.secZenithDelta = parStruct['SECZENITHDELTA'][0]

        if fits_package == 'fitsio':
            self.pwvAtmTable = fitsio.read(self.atmosphereTableFile, ext='PWVATM')
            self.o3AtmTable = fitsio.read(self.atmosphereTableFile, ext='O3ATM')
            self.o2AtmTable = fitsio.read(self.atmosphereTableFile, ext='O2ATM')
            self.rayleighAtmTable = fitsio.read(self.atmosphereTableFile, ext='RAYATM')
        else:
            self.pwvAtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('PWVATM', 1))
            self.o3AtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('O3ATM', 1))
            self.o2AtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('O2ATM', 1))
            self.rayleighAtmTable = pyfits.getdata(self.atmosphereTableFile, ext=('RAYATM', 1))

    def generateTable(self):
        """
        Generate atmosphere table using MODTRAN

        parameters
        ----------
        None

        returns
        -------
        None
        """
        self.atmStd = self.modGen(pmb=self.pmbStd,pwv=self.pwvStd,
                                  o3=self.o3Std,tau=self.tauStd,
                                  alpha=self.alphaStd,zenith=self.zenithStd,
                                  lambdaRange=self.lambdaRange/10.0,
                                  lambdaStep=self.lambdaStep)
        self.atmLambda = self.atmStd['LAMBDA']
        self.atmStdTrans = self.atmStd['COMBINED']

        # get all the steps
        self.pmb = np.linspace(self.atmConfig['pmbRange'][0],
                               self.atmConfig['pmbRange'][1],
                               num=self.atmConfig['pmbSteps'])
        self.pmbDelta = self.pmb[1] - self.pmb[0]
        pmbPlus = np.append(self.pmb, self.pmb[-1] + self.pmbDelta)

        self.lnPwv = np.linspace(np.log(self.atmConfig['pwvRange'][0]),
                                 np.log(self.atmConfig['pwvRange'][1]),
                                 num=self.atmConfig['pwvSteps'])
        self.lnPwvDelta = self.lnPwv[1] - self.lnPwv[0]
        lnPwvPlus = np.append(self.lnPwv, self.lnPwv[-1] + self.lnPwvDelta)
        pwvPlus = np.exp(lnPwvPlus)

        self.o3 = np.linspace(self.atmConfig['o3Range'][0],
                               self.atmConfig['o3Range'][1],
                               num=self.atmConfig['o3Steps'])
        self.o3Delta = self.o3[1] - self.o3[0]
        o3Plus = np.append(self.o3, self.o3[-1] + self.o3Delta)

        self.lnTau = np.linspace(np.log(self.atmConfig['tauRange'][0]),
                                 np.log(self.atmConfig['tauRange'][1]),
                                 num=self.atmConfig['tauSteps'])
        self.lnTauDelta = self.lnTau[1] - self.lnTau[0]
        self.tau = np.exp(self.lnTau)
        lnTauPlus = np.append(self.lnTau, self.lnTau[-1] + self.lnTauDelta)
        tauPlus = np.exp(lnTauPlus)

        self.alpha = np.linspace(self.atmConfig['alphaRange'][0],
                               self.atmConfig['alphaRange'][1],
                               num=self.atmConfig['alphaSteps'])
        self.alphaDelta = self.alpha[1] - self.alpha[0]
        alphaPlus = np.append(self.alpha, self.alpha[-1] + self.alphaDelta)

        self.secZenith = np.linspace(1./np.cos(self.atmConfig['zenithRange'][0]*np.pi/180.),
                                     1./np.cos(self.atmConfig['zenithRange'][1]*np.pi/180.),
                                     num=self.atmConfig['zenithSteps'])
        self.secZenithDelta = self.secZenith[1]-self.secZenith[0]
        self.zenith = np.arccos(1./self.secZenith)*180./np.pi
        secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
        zenithPlus = np.arccos(1./secZenithPlus)*180./np.pi

        # run MODTRAN a bunch of times

        self.fgcmLog.info("Generating %d*%d=%d PWV atmospheres..." % (pwvPlus.size,zenithPlus.size,pwvPlus.size*zenithPlus.size))
        self.pwvAtmTable = np.zeros((pwvPlus.size,zenithPlus.size,self.atmLambda.size))

        for i in range(pwvPlus.size):
            sys.stdout.write('%d' % (i))
            sys.stdout.flush()
            for j in range(zenithPlus.size):
                sys.stdout.write('.')
                sys.stdout.flush()
                atm=self.modGen(pwv=pwvPlus[i],zenith=zenithPlus[j],
                                lambdaRange=self.lambdaRange/10.0,
                                lambdaStep=self.lambdaStep)
                self.pwvAtmTable[i,j,:] = atm['H2O']

        self.fgcmLog.info("\nGenerating %d*%d=%d O3 atmospheres..." % (o3Plus.size,zenithPlus.size,o3Plus.size*zenithPlus.size))
        self.o3AtmTable = np.zeros((o3Plus.size, zenithPlus.size, self.atmLambda.size))

        for i in range(o3Plus.size):
            sys.stdout.write('%d' % (i))
            sys.stdout.flush()
            for j in range(zenithPlus.size):
                sys.stdout.write('.')
                sys.stdout.flush()
                atm=self.modGen(o3=o3Plus[i],zenith=zenithPlus[j],
                                lambdaRange=self.lambdaRange/10.0,
                                lambdaStep=self.lambdaStep)
                self.o3AtmTable[i,j,:] = atm['O3']

        self.fgcmLog.info("\nGenerating %d O2/Rayleigh atmospheres..." % (zenithPlus.size))

        self.o2AtmTable = np.zeros((zenithPlus.size, self.atmLambda.size))
        self.rayleighAtmTable = np.zeros((zenithPlus.size, self.atmLambda.size))

        for j in range(zenithPlus.size):
            sys.stdout.write('.')
            sys.stdout.flush()
            atm=self.modGen(zenith=zenithPlus[j],
                            lambdaRange=self.lambdaRange/10.0,
                            lambdaStep=self.lambdaStep)
            self.o2AtmTable[j,:] = atm['O2']
            self.rayleighAtmTable[j,:] = atm['RAYLEIGH']

        self.fgcmLog.info("\nDone.")

    def saveTable(self, fileName, clobber=False):
        """
        Save atmosphere table to a file

        parameters
        ----------
        fileName: str
           File name to save file
        clobber: bool, default=False
           Clobber fileName if it exists
        """
        # at the moment, only work with fitsio...

        import fitsio

        if os.path.isfile(fileName) and not clobber:
            raise IOError("File %s already exists and clobber is set to False" % (fileName))

        fits = fitsio.FITS(fileName, 'rw')

        parStruct = np.zeros(1, dtype=[('ELEVATION','f4'),
                                       ('PMBELEVATION','f4'),
                                       ('PMBSTD','f4'),
                                       ('PWVSTD','f4'),
                                       ('O3STD','f4'),
                                       ('TAUSTD','f4'),
                                       ('ALPHASTD','f4'),
                                       ('AIRMASSSTD','f4'),
                                       ('LAMBDARANGE','f4',2),
                                       ('LAMBDASTEP','f4'),
                                       ('LAMBDANORM','f4'),
                                       ('ATMLAMBDA','f4',self.atmLambda.size),
                                       ('ATMSTDTRANS','f4',self.atmStdTrans.size),
                                       ('PMB','f4',self.pmb.size),
                                       ('PMBDELTA','f4'),
                                       ('LNPWV','f4',self.lnPwv.size),
                                       ('LNPWVDELTA','f4'),
                                       ('O3','f4',self.o3.size),
                                       ('O3DELTA','f4'),
                                       ('LNTAU','f4',self.lnTau.size),
                                       ('LNTAUDELTA','f4'),
                                       ('ALPHA','f4',self.alpha.size),
                                       ('ALPHADELTA','f4'),
                                       ('SECZENITH','f4',self.secZenith.size),
                                       ('SECZENITHDELTA','f4')])

        parStruct['ELEVATION'] = self.elevation
        parStruct['PMBELEVATION'] = self.pmbElevation
        parStruct['PMBSTD'] = self.pmbStd
        parStruct['PWVSTD'] = self.pwvStd
        parStruct['O3STD'] = self.o3Std
        parStruct['TAUSTD'] = self.tauStd
        parStruct['ALPHASTD'] = self.alphaStd
        parStruct['AIRMASSSTD'] = self.secZenithStd
        parStruct['LAMBDARANGE'][:] = self.lambdaRange
        parStruct['LAMBDASTEP'][:] = self.lambdaStep
        parStruct['LAMBDANORM'][:] = self.lambdaNorm

        parStruct['ATMLAMBDA'][:] = self.atmLambda
        parStruct['ATMSTDTRANS'][:] = self.atmStdTrans

        parStruct['PMB'][:] = self.pmb
        parStruct['PMBDELTA'] = self.pmbDelta
        parStruct['LNPWV'][:] = self.lnPwv
        parStruct['LNPWVDELTA'] = self.lnPwvDelta
        parStruct['O3'][:] = self.o3
        parStruct['O3DELTA'] = self.o3Delta
        parStruct['LNTAU'][:] = self.lnTau
        parStruct['LNTAUDELTA'] = self.lnTauDelta
        parStruct['ALPHA'][:] = self.alpha
        parStruct['ALPHADELTA'] = self.alphaDelta
        parStruct['SECZENITH'][:] = self.secZenith
        parStruct['SECZENITHDELTA'] = self.secZenithDelta

        fits.write_table(parStruct, extname='PARS')

        fits.write_image(self.pwvAtmTable, extname='PWVATM')
        fits.write_image(self.o3AtmTable, extname='O3ATM')
        fits.write_image(self.o2AtmTable, extname='O2ATM')
        fits.write_image(self.rayleighAtmTable, extname='RAYATM')

        fits.close()

    def interpolateAtmosphere(self, lam=None, pmb=None, pwv=None, o3=None, tau=None,
                              alpha=None, zenith=None, ctranslamstd=None):
        """
        Interpolate the atmosphere table to generate an atmosphere without
        requiring MODTRAN be installed.

        parameters
        ----------
        lam: float array, optional
           wavelengths.  Default is self.atmLambda
        pmb: float, optional
           pressure in millibar.  Default to pmbStd
        pwv: float, optional
           pwv in mm.  Default to pwvStd
        o3: float, optional
           ozone in Dobson.  Default to o3Std
        tau: float, optional
           Aerosol optical depth.  Default to tauStd
        alpha: float, optional
           Aerosol slope.  Default to alphaStd
        zenith: float, optional
           Zenith angle (degrees).  Default to zenithStd
        ctranslamstd: [ctrans, lamstd], optional
           Transmission adjustment constant and lambdastd.
           Default to [0.0, lamnorm]

        returns
        -------
        atmInterpolated: float array
           Interpolated atmosphere at self.atmLambda wavelengths
        """

        if lam is None:
            lam = self.atmLambda
        if pmb is None:
            pmb = self.pmbStd
        if pwv is None:
            pwv = self.pwvStd
        if o3 is None:
            o3 = self.o3Std
        if tau is None:
            tau = self.tauStd
        if alpha is None:
            alpha = self.alphaStd
        if zenith is None:
            zenith = self.zenithStd
        if ctranslamstd is None:
            ctranslamstd = [0.0, 7750.0]

        if self.o2Interpolator is None:
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.o2Interpolator = interpolate.RegularGridInterpolator((secZenithPlus, self.atmLambda), self.o2AtmTable)

        if self.rayleighInterpolator is None:
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.rayleighInterpolator = interpolate.RegularGridInterpolator((secZenithPlus, self.atmLambda), self.rayleighAtmTable)

        if self.lnPwvInterpolator is None:
            lnPwvPlus = np.append(self.lnPwv, self.lnPwv[-1] + self.lnPwvDelta)
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.lnPwvInterpolator = interpolate.RegularGridInterpolator((lnPwvPlus, secZenithPlus, self.atmLambda), self.pwvAtmTable)

        if self.o3Interpolator is None:
            o3Plus = np.append(self.o3, self.o3[-1] + self.o3Delta)
            secZenithPlus = np.append(self.secZenith, self.secZenith[-1] + self.secZenithDelta)
            self.o3Interpolator = interpolate.RegularGridInterpolator((o3Plus, secZenithPlus, self.atmLambda), self.o3AtmTable)

        # And finally the pressure correction
        pmbMolecularScattering = np.exp(-(pmb - self.pmbElevation) / self.pmbElevation)
        pmbMolecularAbsorption = pmbMolecularScattering ** 0.6
        pmbFactor = pmbMolecularScattering * pmbMolecularAbsorption

        _secZenith = np.clip(1./np.cos(np.radians(zenith)), self.secZenith[0], self.secZenith[-1])
        _lnPwv = np.clip(np.log(pwv), self.lnPwv[0], self.lnPwv[-1])
        _o3 = np.clip(o3, self.o3[0], self.o3[-1])
        atmInterpolated = (pmbFactor *
                           self.o2Interpolator((_secZenith, lam)) *
                           self.rayleighInterpolator((_secZenith, lam)) *
                           self.lnPwvInterpolator((_lnPwv, _secZenith, lam)) *
                           self.o3Interpolator((_o3, _secZenith, lam)) *
                           (1.0 + ctranslamstd[0] * (lam - ctranslamstd[1]) / ctranslamstd[1]) *
                           np.exp(-1.0 * tau * _secZenith * (lam / self.lambdaNorm)**(-alpha)))

        return atmInterpolated

import numpy,import scipy,import astropy