import os
import shutil
import sys
from collections import defaultdict
import warnings

#Anaconda
import numpy as np
from astropy.time import Time, TimeDelta
from astropy import units as U
from astropy.io import ascii as asc
from astropy.io import fits

#Pip
import extinction
import pymysql
#Local
from .define_filters import define_filters, get_cenwave
from . import connect_to_sndavis, connect_to_supernova
from . import spectroscopy as spec

class LightCurve(object):
    def __init__(self, name, ra, dec, type=None, dist_mod=None, dist_mod_err=None,
                       ebv_mw=None, ebv_mw_err=None, ebv_host=None, ebv_host_err=None,
                       jdexpl=None, jdexplerr=None):
        self.name=name
        self.ra = result['ra0']
        self.dec = result['dec0']
        self.type = result['sntype'] #TODO - update this to look at the type table and make a human readable type
        self.dist_mod = result['mu']
        self.dist_mod_err = result['muerr']
        self.ebv_mw = result['ebvg']
        self.ebv_mw_err = result['ebvgerr']
        self.ebv_host = result['ebvi']
        self.ebv_host_err = result['ebvierr']
        self.jdexpl = result['jdexpl']
        self.jdexpl_err = result['jdexplerr']
        self.phase = defaultdict(list)
        self.jd = defaultdict(list)
        self.apparent_mag = defaultdict(list)
        self.apparent_mag_err = defaultdict(list)
        self.abs_mag = defaultdict(list)
        self.abs_mag_err = defaultdict(list)
        self.A_host = defaultdict(list)
        self.A_err_host = defaultdict(list)
        self.A_mw = defaultdict(list)
        self.A_err_mw = defaultdict(list)

    def add_photometry(filter, jd, phot, phot_err):
        self.jd[filter] = jd
        self.apparent_mag[filter] = phot
        self.apparent_mag_err[filter] = phot_err
        if self.jdexpl:
            self.phase[filter] = self.jd[filter] -self.jdexpl

def calc_abs_mag(app_mag, dist_mod, A_mw, A_host=0, dist_mod_err=0, app_mag_err=0, A_err_mw=0, A_err_host=0):
    '''
    Calculate the absolute magnitude in a given filter given the apparent magnitude, distance modulus, and extinction
    app_mag: arr or float
        array of apparent magnitudes
    dist_mod: float
        distance modulus
    A_mw: float
        Magnitudes of extinction in filter due to the Milky Way
    A_host: float
        Magnitudes of extinction in filter due to host; default=0
    dist_mod_err: float
        error in the distance modulus; default=0
    app_mag_err: arr or float
        array of error values for apparent magnitude (should be same size as app_mag); default=0
    A_err_mw: float
        error on A_mw (in magnitudes); default=0
    A_err_host: float
        err on A_host (in magnitudes); default=0
    '''
    abs_mag = app_mag - dist_mod - A_mw - A_host
    abs_mag_err = np.sqrt(app_mag_err**2 + dist_mod_err**2 + A_err_mw**2 + A_err_host**2)
    return abs_mag, abs_mag_err

class LightCurve2(object):
    '''
    Retrieve information from sndavis database

    Absolute magnitudes are corrected for extinction, apparent magnitude is not
    '''
    def __init__(self, name):
        self.name=name
        self.db, self.cursor = connect_to_sndavis.get_cursor()
        self.id = self.get_sn_id()
        if self.id is not None:
            self.cursor.execute('SELECT * FROM idsupernovae WHERE id = {}'.format(self.id))
            result = self.cursor.fetchone()
            self.ra = result['ra0']
            self.dec = result['dec0']
            self.type = result['sntype'] #TODO - update this to look at the type table and make a human readable type
            self.dist_mod = result['mu']
            self.dist_mod_err = result['muerr']
            self.ebv_mw = result['ebvg']
            self.ebv_mw_err = result['ebvgerr']
            if result['ebvi'] is None:
                self.ebv_host = 0
                self.ebv_host_err = 0
                print('WARNING: host extinciton is None, setting to 0')
            else:
                self.ebv_host = result['ebvi']
                self.ebv_host_err = result['ebvierr']
            self.quality = result['quality']
            self.jdexpl = result['jdexpl']
            self.jdexpl_err = result['jdexplerr']
            self.jd = {}
            self.apparent_mag = {}
            self.apparent_mag_err = {}
            self.phase = {}
            self.abs_mag = {}
            self.abs_mag_err = {}
            self.slopes = {'s1':{}, 's2':{}, 's50':{}, 'tail':{},
                            's1_range':{}, 's2_range':{}, 's50_range':{}, 'tail_range':{},
                            's1_err':{}, 's2_err':{}, 's50_err':{}, 'tail_err':{}}
            self.A_host = {}
            self.A_err_host = {}
            self.A_mw = {}
            self.A_err_mw = {}
            self.jd_limit = {}
            self.phase_limit = {}
            self.apparent_mag_limit = {}
            self.apparent_mag_err_limit = {}
            self.abs_mag_limit = {}
            self.abs_mag_err_limit = {}
            self.telescope = {}
            self.telescope_limit = {}
            self.project = {}
            self.project_limit = {}
        else:
            print('SN {} not found in database'.format(self.name))
            #TODO add an absolute magnitude error

    def get_sn_id(self):
        self.cursor.execute('SELECT targetid FROM supernovanames WHERE name = {}'.format('"{}"'.format(self.name)))
        result = self.cursor.fetchone()
        if result is None:
            return None
        return result['targetid']

    def get_photometry(self, band='all'):
        '''
        Get apparent magnitude of a specific band
        '''
        filter_dict =  define_filters()
        if band == 'all':
            self.cursor.execute("SELECT DISTINCT BINARY(filter) FROM photometry WHERE targetid={}".format(self.id))
            results = self.cursor.fetchall()
            bands = [iband['BINARY(filter)'].decode('utf-8') for iband in results]

        elif len(band)==1 or isinstance(band, str): #Single filter
            bands = [band]
        else: #multiple filters
            bands = band
        for ifilter in bands:
            if ifilter in filter_dict.keys():
                A_host_band, A_err_host_band = spec.calc_extinction(self.ebv_host, ifilter)
                A_mw_band, A_err_mw_band = spec.calc_extinction(self.ebv_mw, ifilter)
            else:
                print('Could not calculate extinction for {} because not in filter dictionary'.format(ifilter))
                A_host_band = 0
                A_err_host_band = 0
                A_mw_band = 0
                A_err_mw_band = 0
            #self.cursor.execute("SELECT jd, mag, magerr, datatype, configuration  FROM photometry JOIN photometrysource on photometry.source=photometrysource.source WHERE targetid={} AND filter=BINARY('{}')".format(self.id, ifilter))
            self.cursor.execute("SELECT jd, mag, magerr, datatype, source FROM photometry WHERE targetid={} AND filter=BINARY('{}')".format(self.id, ifilter))
            results = self.cursor.fetchall()
            jd = []
            mag = []
            mag_err = []
            telescope = []
            project=[]
            jd_limit = []
            mag_limit = []
            mag_err_limit = []
            telescope_limit = []
            project_limit = []
            for irow in results:
                if irow['magerr'] is None:
                    irow['magerr'] = 0
                if irow['datatype'] < 0: #photometry represents a limit
                    jd_limit.append(float(irow['jd']))
                    mag_limit.append(float(irow['mag']))
                    mag_err_limit.append(float(irow['magerr']))
                    #telescope_limit.append(irow['configuration'])
                    #telescope_limit.append(irow['source'])
                    #if irow['configuration'].startswith('1m'):
                    #    project_limit.append('LCO')
                    #elif irow['configuration'] == 'Meckering' or irow['configuration'] == 'Prompt5':
                    #    project_limit.append('DLT40')
                    #else:
                    #    project_limit.append(None)
                    telescope_limit.append(irow['source'])
                    if irow['source'] in [5209, 5109, 5009, 4809, 4909, 4709, 4609, 4509, 140, 141, 142, 165, 166, 167]:
                        project_limit.append('LCO')
                    elif irow['source'] in [5300, 5400]:
                        project_limit.append('DLT40')
                    else:
                        project_limit.append(None)
                else:
                    jd.append(float(irow['jd']))
                    mag.append(float(irow['mag']))
                    mag_err.append(float(irow['magerr']))
                    # telescope.append(irow['configuration'])
                    # if irow['configuration'].startswith('1m'):
                    #     project.append('LCO')
                    # elif irow['configuration'] == 'Meckering' or irow['configuration'] == 'Prompt5':
                    #     project.append('DLT40')
                    # else:
                    #     project.append(None)
                    telescope.append(irow['source'])
                    if irow['source'] in [5209, 5109, 5009, 4809, 4909, 4709, 4609, 4509, 140, 141, 142, 165, 166, 167]:
                        project.append('LCO')
                    elif irow['source'] in [5300, 5400]:
                        project.append('DLT40')
                    else:
                        project.append(None)
            self.jd[ifilter] = np.array(jd)
            self.apparent_mag[ifilter] = np.array(mag)
            self.apparent_mag_err[ifilter] = np.array(mag_err)
            if self.jdexpl is not None:
                self.phase[ifilter] = self.jd[ifilter] - self.jdexpl
            else:
                print('WARNING: No explosion epoch in database, phase is not calculated')
            self.A_host[ifilter] = A_host_band
            self.A_err_host[ifilter] = A_err_host_band
            self.A_mw[ifilter] = A_mw_band
            self.A_err_mw[ifilter] = A_err_mw_band
            self.jd_limit[ifilter] = np.array(jd_limit)
            self.apparent_mag_limit[ifilter] = np.array(mag_limit)
            self.apparent_mag_err_limit[ifilter] = np.array(mag_err_limit)
            if self.jdexpl is not None:
                self.phase_limit[ifilter] = self.jd_limit[ifilter] - self.jdexpl
            else:
                print('WARNING: No explosion epoch in database, phase is not calculated')
            self.telescope[ifilter] = np.array(telescope)
            self.telescope_limit[ifilter] = np.array(telescope_limit)
            self.project[ifilter] = np.array(project)
            self.project_limit[ifilter] = np.array(project_limit)

    def get_abs_mag(self, band='all'):
        '''
        Calcualte absolute magnitude
        '''
        print('Calculating Absolute Magntidue with Extinction')
        if band == 'all':
            bands = self.apparent_mag.keys()
        elif len(band)==1:
            bands = [band]
        else:
            bands = band
        if self.dist_mod_err is None:
            self.dist_mod_err = 0
        for iband in bands:
            if self.ebv_host is None:
                self.abs_mag[iband], self.abs_mag_err[iband] = calc_abs_mag(self.apparent_mag[iband],
                                                   self.dist_mod,
                                                   self.A_mw[iband],
                                                   app_mag_err=self.apparent_mag_err[iband],
                                                   dist_mod_err=self.dist_mod_err)
                self.abs_mag_limit[iband], self.abs_mag_err_limit[iband] = calc_abs_mag(self.apparent_mag_limit[iband],
                                                   self.dist_mod,
                                                   self.A_mw[iband],
                                                   app_mag_err=self.apparent_mag_err_limit[iband],
                                                   dist_mod_err=self.dist_mod_err)
            else:
                self.abs_mag[iband], self.abs_mag_err[iband] = calc_abs_mag(self.apparent_mag[iband],
                                                                            self.dist_mod,
                                                                            self.A_mw[iband],
                                                                            A_host=self.A_host[iband],
                                                                            app_mag_err=self.apparent_mag_err[iband],
                                                                            dist_mod_err=self.dist_mod_err,
                                                                            A_err_mw=self.A_err_mw[iband],
                                                                            A_err_host=self.A_err_host[iband])
                self.abs_mag_limit[iband], self.abs_mag_err_limit[iband] = calc_abs_mag(self.apparent_mag_limit[iband],
                                                                            self.dist_mod,
                                                                            self.A_mw[iband],
                                                                            A_host=self.A_host[iband],
                                                                            app_mag_err=self.apparent_mag_err_limit[iband],
                                                                            dist_mod_err=self.dist_mod_err,
                                                                            A_err_mw=self.A_err_mw[iband],
                                                                            A_err_host=self.A_err_host[iband])

    def get_slope(self,slope_type, band='V'):
        '''
        '''
        self.cursor.execute('SELECT * FROM snslope WHERE  targetid={} AND \
                                                         slopetype="{}" AND \
                                                         filter=BINARY("{}")'.format(self.id, slope_type, band))
        result = self.cursor.fetchone()
        if result is not None:
            self.slopes['{}'.format(slope_type)][band]       = result['slope']
            self.slopes['{}_err'.format(slope_type)][band]   = result['slopeerr']
            self.slopes['{}_range'.format(slope_type)][band] = (result['tstart'], result['tstop'])
        else:
            self.slopes['{}'.format(slope_type)][band]       = None
            self.slopes['{}_err'.format(slope_type)][band]   = None
            self.slopes['{}_range'.format(slope_type)][band] = [None, None]

########################################

class LightCurveLCO(object):
    '''
    Retrieve information from LCO supernova database
    '''
    def __init__(self, name, **kwargs):
        if 'config_dir' not in kwargs.keys():
            kwargs['config_dir'] = os.environ['LCOSNDIR']
        if 'config_filename' not in kwargs.keys():
            kwargs['config_filename'] = 'configure'
        self.name=name
        self.db, self.cursor = connect_to_supernova.get_cursor(**kwargs)
        self.id = self.get_sn_id()
        if self.id is not None:
            self.cursor.execute('SELECT * FROM targets WHERE id = {}'.format(self.id))
            result = self.cursor.fetchone()
            self.ra = result['ra0']
            self.dec = result['dec0']
            self.type = result['classification'] #TODO - update this to look at the type table and make a human readable type
            self.z = result['redshift']
            #These are set by the get_photometry method
            self.quality = {}
            self.jd = {}
            self.apparent_mag = {}
            self.apparent_mag_err = {}
            self.telescope = {}

    def get_sn_id(self):
        num_results = self.cursor.execute("SELECT DISTINCT targetid FROM targetnames WHERE name LIKE '%{}%'".format(self.name))
        result = self.cursor.fetchall()
        if num_results == 0:
            return None
        elif num_results>1:
            warnings.warn('Multiple targetids match your query. Please enter the targetid you would like to select')
            for iresult in result:
                target = self.cursor.execute("SELECT name FROM targetnames WHERE targetid = {}".format(iresult['targetid']))
                target_results = target.fetchall()
                print('targetid: {}'.format(iresult['targetid']))
                for itarget_result in target_results:
                    print(itarget_result['name'])
            targetid = input('Enter target id: ')
        else:
            targetid = result[0]['targetid']
        return targetid

    def get_photometry(self, band='all'):
        '''
        Get apparent magnitude of a specific band
        '''
        filter_dict =  define_filters()
        if band == 'all':
            self.cursor.execute("SELECT DISTINCT BINARY(filter) FROM photlco WHERE targetid={}".format(self.id))
            results = self.cursor.fetchall()
            bands = [iband['BINARY(filter)'].decode('utf-8') for iband in results]

        elif len(band)==1 or isinstance(band, str): #Single filter
            bands = [band]
        else: #multiple filters
            bands = band
        for ifilter in bands:
            #self.cursor.execute("SELECT jd, mag, magerr, datatype, configuration  FROM photometry JOIN photometrysource on photometry.source=photometrysource.source WHERE targetid={} AND filter=BINARY('{}')".format(self.id, ifilter))
            sql_str = "SELECT mjd, mag, dmag, telescope, quality FROM photlco WHERE"\
                                + " wcs != 9999 AND (psf != 'X' AND psf != 9999) AND (zcat != 'X' OR abscat !='X') AND mag != 9999"\
                                + " AND targetid={} AND filter=BINARY('{}')".format(self.id, ifilter)
            self.cursor.execute(sql_str)
            results = self.cursor.fetchall()
            jd = []
            mag = []
            mag_err = []
            quality = []
            telescope = []
            for irow in results:
                if irow['dmag'] is None:
                    irow['dmag'] = 0
                jd.append(Time(float(irow['mjd']), format='mjd').jd)
                mag.append(float(irow['mag']))
                mag_err.append(float(irow['dmag']))
                quality.append(int(irow['quality']))
                telescope.append(irow['telescope'])
            self.jd[ifilter] = np.array(jd)
            self.apparent_mag[ifilter] = np.array(mag)
            self.apparent_mag_err[ifilter] = np.array(mag_err)
            self.quality[ifilter] = np.array(quality)
            self.telescope[ifilter] = np.array(telescope)
################

def get_closest_photometry(date_obs, jd, phot_mag):
    '''
    Find the closest photometric observations to a given date

    Inputs:
    ----------
        date_obs: Time object
            astropy Time object of the observervation date that you want to find the closest observations to
        jd: list
            list of photometric observations in JD
        phot_mag: list
            list of magnitudes corresponding to the jd list for a single filter

    Returns:
    -----------
        closest_mag: float
            the closest magnitude to the date_obs given
        date_sep: float
            the difference in days between the observation and the closest photometric point
    '''
    date_indx = np.argmin(np.abs(jd - date_obs.jd))
    closest_mag = phot_mag[date_indx]
    date_sep = (jd - date_obs.jd)[date_indx]
    return closest_mag, date_sep

def get_interpolated_photometry(date_obs, jd, phot_mag):
    '''
    Interpolate between photometric points in a given band to get the value at a single date
        Inputs:
    ----------
        date_obs: Time object
            astropy Time object of the observervation date that you want to find the closest observations to
        jd: list
            list of photometric observations in JD
        phot_mag: list
            list of magnitudes corresponding to the jd list for a single filter

    Returns:
    -----------
        interp_mag: float
            the magnitude interpolated to the date_obs given
        date_sep: tup
            the difference in days between the observation and the two points used in the interpolation
    '''

    before_indx = jd < date_obs.jd
    after_indx = jd >= date_obs.jd
    if (before_indx==False).all():
        before_sep = None
    else:
        before_sep = Time(jd[before_indx][-1], format='jd') - date_obs
    if (after_indx == False).all():
        after_sep = None
    else:
        after_sep =  Time(jd[after_indx][0],   format='jd') - date_obs
    if (before_sep is not None) and (after_sep is not None):
        interp_mag = np.interp(date_obs.jd, jd, phot_mag)
    else:
        interp_mag = None
    return interp_mag, (before_sep, after_sep)

def convert_mag_to_flux(phot_mag, ifilter):
    '''
    Converts magnitude units to flux units (ergs/cm^2/s/A)

    Inputs:
    --------
    phot_mag: list or array
        list or array of magnitudes
    ifilter: str
        filter that phot_mag observations were taken with

    Returns:
    --------
    flux : array
        list of flux values
    '''
    bandpar_dict = define_filters()
    flux = 10**(np.array(phot_mag)/-2.5)*bandpar_dict[ifilter][4]
    return flux

def get_cenwave(ifilter):
    '''
    return the central wavelength of a given filter

    Inputs:
    -------
    ifilter: str
        the name of a filter (must be present in define_filters)

    Returns:
    --------
    cenwave: float
        central wavelength of filter
    '''
    bandpar_dict = define_filters()
    cenwave = bandpar_dict[ifilter][2]
    return cenwave

def scale_spectra_quba(snname, filename, filter_dir=None, date_kw='date-obs', max_sep=7,
                        header_date=False, sndavis=True, lightcurve=None, verbose=False):
    '''
    You must be connected to dark
    filter_dir should be /Users/bostroem/Dropbox/DLT40_all/script/scalespectra/filters/
    if header_date is True, then date_kw contains the keyword to read from the header.
    if header_date is False, then the date_kw should be treated as the date of the observation

    Inputs:
    ---------
    snname: str
        name of supernova - will be used to find in SNDAVIS database is sndavis is True
    filename: str
        name of spectrum file.
    filter_dir: str
        location of information filter throughput
    date_kw: str
        if header_date is True: name of keyword in fits file to use to find date
        if header_date is False: date of observation
    max_set: float
        number of days that can separate photometry point and date of spectrum
    header_date: bool
        controls whether spectrum date is read from header
    sndavis: bool
        if True: use SNDAVIS database to get light curve
        if False: pass a light curve into the light curve keyword
    lightcurve: LightCurve object
        an object that has a jd attribute and an app_mag attribute that are dictionaries with filters as keys

    '''
    import qubascalespectra
    if lightcurve is None:
        lightcurve = LightCurve2(snname)
        lightcurve.get_photometry()
    if header_date is True:
        date_obs = Time(fits.getval(filename, date_kw, 0))
    else:
        date_obs = Time(date_kw)
    spec_phase = (date_obs - Time(lightcurve.jdexpl, format='jd')).value
    band = ''
    mphot = []
    for ifilter in lightcurve.jd.keys():
        if ifilter not in ['us', 'vs', 'bs', 'uw1', 'uw2', 'um1', 'um2', 'UVW1', 'UVW2', 'UVM1', 'UVM2']: #avoid swift filters
            cenwave = get_cenwave(ifilter)
            if (cenwave > 3500):
                mag, date_sep = get_interpolated_photometry(date_obs, lightcurve.jd[ifilter], lightcurve.apparent_mag[ifilter])
                if mag is not None:
                    print('using {}={} interpolated to {} from {} and {}'.format(ifilter, mag, date_obs.iso,
                                                                               (date_sep[0]+date_obs).iso,
                                                                               (date_sep[1]+date_obs).iso))
                    if ifilter.endswith('p'): #strip off p for LCO filters since qubascalespectra assumes filters are single character
                        band = band+ifilter[0]
                    else:
                        band=band+ifilter
                    mphot.append(mag)
                else:
                    if date_sep[0] is None:
                        print('No data for {} before {}'.format(ifilter, date_obs))
                    if date_sep[1] is None:
                        print('No data for {} after {}'.format(ifilter, date_obs))
    if len(mphot) > 0:
        if verbose:
            print('bands: {}, photometry: {}'.format(band, mphot))
        qubascalespectra.scale_spectrum(filename, band, mphot, filter_dir)
        shutil.move('log.txt', filename.split('.')[0]+'.log')
        with open(filename.split('.')[0]+'.log', 'a') as ofile:
            ofile.write(band+'\n')
            mphot_str = [str(iphot) for iphot in mphot]
            ofile.write(','.join(mphot_str)+'\n')
    else:
        print('WARNING: no photometry within {} days of observation date ({}) ({})'.format(max_sep, date_obs, filename))


""" Parse q-chem log files and created displaced structures
    using normal modes
"""

import numpy as np
import glob
import sys
import cclib
from cclib.parser import ccopen

NAME = {
    1:  'H'  ,
    2:  'He' ,
    3:  'Li' ,
    4:  'Be' ,
    5:  'B'  ,
    6:  'C'  ,
    7:  'N'  ,
    8:  'O'  ,
    9:  'F'  ,
   10:  'Ne' ,
   11:  'Na' ,
   12:  'Mg' ,
   13:  'Al' ,
   14:  'Si' ,
   15:  'P'  ,
   16:  'S'  ,
   17:  'Cl' ,
   18:  'Ar' ,
   19:  'K'  ,
   20:  'Ca' ,
   21:  'Sc' ,
   22:  'Ti' ,
   23:  'V'  ,
   24:  'Cr' ,
   25:  'Mn' ,
   26:  'Fe' ,
   27:  'Co' ,
   28:  'Ni' ,
   29:  'Cu' ,
   30:  'Zn' ,
   31:  'Ga' ,
   32:  'Ge' ,
   33:  'As' ,
   34:  'Se' ,
   35:  'Br' ,
   36:  'Kr' ,
   37:  'Rb' ,
   38:  'Sr' ,
   39:  'Y'  ,
   40:  'Zr' ,
   41:  'Nb' ,
   42:  'Mo' ,
   43:  'Tc' ,
   44:  'Ru' ,
   45:  'Rh' ,
   46:  'Pd' ,
   47:  'Ag' ,
   48:  'Cd' ,
   49:  'In' ,
   50:  'Sn' ,
   51:  'Sb' ,
   52:  'Te' ,
   53:  'I'  ,
   54:  'Xe' ,
   55:  'Cs' ,
   56:  'Ba' ,
   57:  'La' ,
   58:  'Ce' ,
   59:  'Pr' ,
   60:  'Nd' ,
   61:  'Pm' ,
   62:  'Sm' ,
   63:  'Eu' ,
   64:  'Gd' ,
   65:  'Tb' ,
   66:  'Dy' ,
   67:  'Ho' ,
   68:  'Er' ,
   69:  'Tm' ,
   70:  'Yb' ,
   71:  'Lu' ,
   72:  'Hf' ,
   73:  'Ta' ,
   74:  'W'  ,
   75:  'Re' ,
   76:  'Os' ,
   77:  'Ir' ,
   78:  'Pt' ,
   79:  'Au' ,
   80:  'Hg' ,
   81:  'Tl' ,
   82:  'Pb' ,
   83:  'Bi' ,
   84:  'Po' ,
   85:  'At' ,
   86:  'Rn' ,
   87:  'Fr' ,
   88:  'Ra' ,
   89:  'Ac' ,
   90:  'Th' ,
   91:  'Pa' ,
   92:  'U'  ,
   93:  'Np' ,
   94:  'Pu' ,
   95:  'Am' ,
   96:  'Cm' ,
   97:  'Bk' ,
   98:  'Cf' ,
   99:  'Es' ,
  100:  'Fm' ,
  101:  'Md' ,
  102:  'No' ,
  103:  'Lr' ,
  104:  'Rf' ,
  105:  'Db' ,
  106:  'Sg' ,
  107:  'Bh' ,
  108:  'Hs' ,
  109:  'Mt' ,
  110:  'Ds' ,
  111:  'Rg' ,
  112:  'Cn' ,
  114:  'Uuq',
  116:  'Uuh'}

def create_displaced_structures():
    """ Creates 10 displaced xyz files for T = 300, 600, 1200
        by normal mode sampling.
    """
    log_files = glob.glob("./log_files/**/*.log", recursive=True)
    output_folder = "./xyz"
    # Makes flat modes less likely to explode molecule
    force_min = 0.2
    for i, log_file in enumerate(log_files):
        if i < 43000:
            continue
        if i % 500 == 0:
            print(f"{i} of {len(log_files)} processed")
        mylogfile = ccopen(log_file)
        data = mylogfile.parse()
        force = get_force_constants(log_file)
        coords = data.atomcoords[-1]
        atoms = [NAME[i] for i in data.atomnos]
        modes = data.vibdisps
        # Only keep modes from last job
        modes = modes[-force.size:]
        normalized_modes = modes / np.linalg.norm(modes, axis=(1,2))[:,None,None]
        write_xyz(f"{output_folder}/{i}_0.xyz", coords, atoms)
        counter = 0
        for T in 300, 600, 1200:
            for n in range(10):
                dcoords = coords[:]
                counter += 1
                c = np.random.uniform(0, 1, size=len(force))
                c /= np.sum(c)
                r_scale = np.random.uniform(0, 1)
                c *= r_scale
                for j, mode in enumerate(normalized_modes):
                    sign = np.random.choice([-1.0, 1.0])
                    frc = max(force_min, force[j])
                    r = sign * np.sqrt(3 * c[j] * len(atoms) * 1.380e-5 * T / frc)
                    dcoords += mode * r
                write_xyz(f"{output_folder}/{i}_{counter}.xyz", dcoords, atoms)

def write_xyz(filename, coords, atoms):
    """ Writes xyz files
    """
    with open(filename, "w") as f:
        f.write(f"{len(atoms)}\n")
        for j, atom in enumerate(atoms):
            f.write("\n%2s %20.12f %20.12f %20.12f" % \
                (atom, coords[j,0], coords[j,1], coords[j,2]))
#
# atomcoords, atommasses, charge, mult, vibdisps, vibfreqs

def get_force_constants(filename):
    """ Reads the force constants of the last job in
        the Q-Chem output file
    """
    force_constants = []
    last_job_flag = False
    with open(filename) as f:
        for line in f:
            if line.startswith("Running Job"):
                tokens = line.split()
                current_job = tokens[2]
                last_job = tokens[4]
                if current_job == last_job:
                    last_job_flag = True
            if last_job_flag and line.startswith(" Force Cnst:"):
                tokens = line.split()
                force_constants.extend(tokens[2:])
            else:
                continue
    #try:
    return np.asarray(force_constants, dtype=float)
    #except:
    #    print(filename)
    #    print(force_constants)
    #    quit()

if __name__ == "__main__":
    create_displaced_structures()


# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Oliver J. Backhouse <olbackhouse@gmail.com>
#         George H. Booth <george.booth@kcl.ac.uk>
#

'''
Auxiliary second-order Green's function perturbation theory
with density fitting
'''

import numpy as np
import ctypes
from pyscf import lib
from pyscf.lib import logger
from pyscf import __config__
from pyscf import ao2mo, df
from pyscf.agf2 import ragf2, mpi_helper, _agf2
from pyscf.agf2 import aux_space as aux

BLKMIN = getattr(__config__, 'agf2_blkmin', 100)

def build_se_part(agf2, eri, gf_occ, gf_vir, os_factor=1.0, ss_factor=1.0):
    ''' Builds either the auxiliaries of the occupied self-energy,
        or virtual if :attr:`gf_occ` and :attr:`gf_vir` are swapped.

    Args:
        eri : _ChemistsERIs
            Electronic repulsion integrals
        gf_occ : GreensFunction
            Occupied Green's function
        gf_vir : GreensFunction
            Virtual Green's function

    Kwargs:
        os_factor : float
            Opposite-spin factor for spin-component-scaled (SCS)
            calculations. Default 1.0
        ss_factor : float
            Same-spin factor for spin-component-scaled (SCS)
            calculations. Default 1.0

    Returns:
        :class:`SelfEnergy`
    '''

    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(agf2.stdout, agf2.verbose)

    assert type(gf_occ) is aux.GreensFunction
    assert type(gf_vir) is aux.GreensFunction

    nmo = eri.nmo
    nocc, nvir = gf_occ.naux, gf_vir.naux
    naux = agf2.with_df.get_naoaux()
    tol = agf2.weight_tol
    facs = dict(os_factor=os_factor, ss_factor=ss_factor)

    ei, ci = gf_occ.energy, gf_occ.coupling
    ea, ca = gf_vir.energy, gf_vir.coupling

    qxi, qja = _make_qmo_eris_incore(agf2, eri, (ci, ci, ca))

    himem_required = naux*(nvir+nmo) + (nocc*nvir)*(2*nmo+1) + (2*nmo**2)
    himem_required *= 8e-6
    himem_required *= lib.num_threads()

    if ((himem_required*1.05 + lib.current_memory()[0]) > agf2.max_memory
            and agf2.allow_lowmem_build) or agf2.allow_lowmem_build == 'force':
        log.debug('Thread-private memory overhead %.3f exceeds max_memory, using '
                  'low-memory version.', himem_required)
        vv, vev = _agf2.build_mats_dfragf2_lowmem(qxi, qja, ei, ea, **facs)
    else:
        vv, vev = _agf2.build_mats_dfragf2_incore(qxi, qja, ei, ea, **facs)

    e, c = _agf2.cholesky_build(vv, vev)
    se = aux.SelfEnergy(e, c, chempot=gf_occ.chempot)
    se.remove_uncoupled(tol=tol)

    if not (agf2.frozen is None or agf2.frozen == 0):
        mask = ragf2.get_frozen_mask(agf2)
        coupling = np.zeros((nmo, se.naux))
        coupling[mask] = se.coupling
        se = aux.SelfEnergy(se.energy, coupling, chempot=se.chempot)

    log.timer('se part', *cput0)

    return se

def get_jk(agf2, eri, rdm1, with_j=True, with_k=True):
    ''' Get the J/K matrices.

    Args:
        eri : ndarray or H5 dataset
            Electronic repulsion integrals (NOT as _ChemistsERIs). In
            the case of no bra/ket symmetry, a tuple can be passed.
        rdm1 : 2D array
            Reduced density matrix

    Kwargs:
        with_j : bool
            Whether to compute J. Default value is True
        with_k : bool
            Whether to compute K. Default value is True

    Returns:
        tuple of ndarrays corresponding to J and K, if either are
        not requested then they are set to None.
    '''

    nmo = rdm1.shape[0]
    npair = nmo*(nmo+1)//2
    naux = agf2.with_df.get_naoaux()
    vj = vk = None

    if with_j:
        rdm1_tril = lib.pack_tril(rdm1 + np.tril(rdm1, k=-1))
        vj = np.zeros((npair,))

    if with_k:
        vk = np.zeros((nmo, nmo))

    fdrv = ao2mo._ao2mo.libao2mo.AO2MOnr_e2_drv
    fmmm = ao2mo._ao2mo.libao2mo.AO2MOmmm_bra_nr_s2
    ftrans = ao2mo._ao2mo.libao2mo.AO2MOtranse2_nr_s2

    if isinstance(eri, tuple):
        bra, ket = eri
    else:
        bra = ket = eri

    blksize = _agf2.get_blksize(agf2.max_memory, (npair, npair, 1, nmo**2, nmo**2))
    blksize = min(nmo, max(BLKMIN, blksize))
    logger.debug1(agf2, 'blksize (dfragf2.get_jk) = %d' % blksize)
    buf = (np.empty((blksize, nmo, nmo)), np.empty((blksize, nmo, nmo)))

    for p0, p1 in mpi_helper.prange(0, naux, blksize):
        bra0 = bra[p0:p1]
        ket0 = ket[p0:p1]
        rho = np.dot(ket0, rdm1_tril)

        if with_j:
            vj += np.dot(rho, bra0)

        if with_k:
            buf1 = buf[0][:p1-p0]
            fdrv(ftrans, fmmm,
                 buf1.ctypes.data_as(ctypes.c_void_p),
                 bra0.ctypes.data_as(ctypes.c_void_p),
                 rdm1.ctypes.data_as(ctypes.c_void_p),
                 ctypes.c_int(p1-p0), ctypes.c_int(nmo),
                 (ctypes.c_int*4)(0, nmo, 0, nmo),
                 lib.c_null_ptr(), ctypes.c_int(0))

            buf2 = lib.unpack_tril(ket0, out=buf[1])
            buf1 = buf1.reshape(-1, nmo)
            buf2 = buf2.reshape(-1, nmo)

            vk = lib.dot(buf1.T, buf2, c=vk, beta=1)

    if with_j:
        mpi_helper.barrier()
        mpi_helper.allreduce_safe_inplace(vj)
        mpi_helper.barrier()
        vj = lib.unpack_tril(vj)

    if with_k:
        mpi_helper.barrier()
        mpi_helper.allreduce_safe_inplace(vk)

    return vj, vk

class DFRAGF2(ragf2.RAGF2):
    ''' Restricted AGF2 with canonical HF reference with density fitting

    Attributes:
        verbose : int
            Print level. Default value equals to :class:`Mole.verbose`
        max_memory : float or int
            Allowed memory in MB. Default value equals to :class:`Mole.max_memory`
        incore_complete : bool
            Avoid all I/O. Default is False.
        allow_lowmem_build : bool or str
            Allow the self-energy build to switch to a serially slower
            code with lower thread-private memory overhead if needed. One
            of True, False or 'force'. Default value is True.
        conv_tol : float
            Convergence threshold for AGF2 energy. Default value is 1e-7
        conv_tol_rdm1 : float
            Convergence threshold for first-order reduced density matrix.
            Default value is 1e-8.
        conv_tol_nelec : float
            Convergence threshold for the number of electrons. Default
            value is 1e-6.
        max_cycle : int
            Maximum number of AGF2 iterations. Default value is 50.
        max_cycle_outer : int
            Maximum number of outer Fock loop iterations. Default
            value is 20.
        max_cycle_inner : int
            Maximum number of inner Fock loop iterations. Default
            value is 50.
        weight_tol : float
            Threshold in spectral weight of auxiliaries to be considered
            zero. Default 1e-11.
        diis : bool or lib.diis.DIIS
            Whether to use DIIS, can also be a lib.diis.DIIS object. Default
            value is True.
        diis_space : int
            DIIS space size. Default value is 8.
        diis_min_space : int
            Minimum space of DIIS. Default value is 1.
        fock_diis_space : int
            DIIS space size for Fock loop iterations. Default value is 6.
        fock_diis_min_space :
            Minimum space of DIIS. Default value is 1.
        os_factor : float
            Opposite-spin factor for spin-component-scaled (SCS)
            calculations. Default 1.0
        ss_factor : float
            Same-spin factor for spin-component-scaled (SCS)
            calculations. Default 1.0
        damping : float
            Damping factor for the self-energy. Default value is 0.0

    Saved results

        e_corr : float
            AGF2 correlation energy
        e_tot : float
            Total energy (HF + correlation)
        e_1b : float
            One-body part of :attr:`e_tot`
        e_2b : float
            Two-body part of :attr:`e_tot`
        e_init : float
            Initial correlation energy (truncated MP2)
        converged : bool
            Whether convergence was successful
        se : SelfEnergy
            Auxiliaries of the self-energy
        gf : GreensFunction
            Auxiliaries of the Green's function
    '''

    def __init__(self, mf, frozen=None, mo_energy=None, mo_coeff=None, mo_occ=None):
        ragf2.RAGF2.__init__(self, mf, frozen=frozen, mo_energy=mo_energy,
                             mo_coeff=mo_coeff, mo_occ=mo_occ)

        if getattr(mf, 'with_df', None) is not None:
            self.with_df = mf.with_df
        else:
            self.with_df = df.DF(mf.mol)
            self.with_df.auxbasis = df.make_auxbasis(mf.mol, mp2fit=True)

        self.allow_lowmem_build = True

        self._keys.update(['_with_df', 'allow_lowmem_build'])

    build_se_part = build_se_part
    get_jk = get_jk

    def ao2mo(self, mo_coeff=None):
        ''' Get the density-fitted electronic repulsion integrals in
            MO basis.
        '''

        eri = _make_mo_eris_incore(self, mo_coeff)

        return eri

    def reset(self, mol=None):
        self.with_df.reset(mol)
        return ragf2.RAGF2.reset(self, mol)

    @property
    def with_df(self):
        return self._with_df
    @with_df.setter
    def with_df(self, val):
        self._with_df = val
        self._with_df.__class__ = DF

class DF(df.DF):
    ''' Replaces the :class:`DF.prange` function with one which
        natively supports MPI, if used.
    '''
    def prange(self, start=None, stop=None, step=None):
        if start is None: start = 0
        if stop is None: stop = self.get_naoaux()
        if step is None: step = self.blockdim

        for p0, p1 in mpi_helper.prange(start, stop, step):
            yield p0, p1

class _ChemistsERIs(ragf2._ChemistsERIs):
    ''' (pq|rs) as (pq|J)(J|rs)

    MO tensors stored in tril form, we only need QMO tensors
    in low-symmetry
    '''
    pass

def _make_mo_eris_incore(agf2, mo_coeff=None):
    ''' Returns _ChemistsERIs
    '''

    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(agf2.stdout, agf2.verbose)

    eris = _ChemistsERIs()
    eris._common_init_(agf2, mo_coeff)
    with_df = agf2.with_df
    nmo = eris.fock.shape[0]
    npair = nmo*(nmo+1)//2
    naux = with_df.get_naoaux()

    qxy = np.zeros((naux, npair))
    mo = np.asarray(eris.mo_coeff, order='F')
    sij = (0, nmo, 0, nmo)
    sym = dict(aosym='s2', mosym='s2')

    for p0, p1 in with_df.prange():
        eri0 = with_df._cderi[p0:p1]
        qxy[p0:p1] = ao2mo._ao2mo.nr_e2(eri0, mo, sij, out=qxy[p0:p1], **sym)

    mpi_helper.barrier()
    mpi_helper.allreduce_safe_inplace(qxy)

    eris.eri = eris.qxy = qxy

    log.timer('MO integral transformation', *cput0)

    return eris

def _make_qmo_eris_incore(agf2, eri, coeffs):
    ''' Returns tuple of ndarray
    '''

    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(agf2.stdout, agf2.verbose)

    cx = np.eye(agf2.nmo)
    if not (agf2.frozen is None or agf2.frozen == 0):
        mask = ragf2.get_frozen_mask(agf2)
        cx = cx[:,mask]

    nmo = eri.fock.shape[0]
    npair = nmo*(nmo+1)//2
    with_df = agf2.with_df
    naux = with_df.get_naoaux()
    ci, cj, ca = coeffs

    xisym, nxi, cxi, sxi = ao2mo.incore._conc_mos(cx, ci, compact=False)
    jasym, nja, cja, sja = ao2mo.incore._conc_mos(cj, ca, compact=False)
    sym = dict(aosym='s2', mosym='s1')

    qxi = np.zeros((naux, nxi))
    qja = np.zeros((naux, nja))
    buf = np.zeros((with_df.blockdim, npair))

    for p0, p1 in mpi_helper.prange(0, naux, with_df.blockdim):
        naux0 = p1 - p0
        buf0 = buf[:naux0]
        buf0[:] = eri.eri[p0:p1]

        qxi[p0:p1] = ao2mo._ao2mo.nr_e2(buf0, cxi, sxi, out=qxi[p0:p1], **sym)
        qja[p0:p1] = ao2mo._ao2mo.nr_e2(buf0, cja, sja, out=qja[p0:p1], **sym)

    qxi = qxi.reshape(naux, -1)
    qja = qja.reshape(naux, -1)

    mpi_helper.barrier()
    mpi_helper.allreduce_safe_inplace(qxi)
    mpi_helper.allreduce_safe_inplace(qja)

    log.timer('QMO integral transformation', *cput0)

    return (qxi, qja)

if __name__ == '__main__':
    from pyscf import gto, scf, mp

    mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='cc-pvdz', verbose=3)
    rhf = scf.RHF(mol).density_fit()
    rhf.conv_tol = 1e-11
    rhf.run()

    ragf2 = DFRAGF2(rhf)

    ragf2.run()
    ragf2.ipagf2(nroots=5)
    ragf2.eaagf2(nroots=5)


__all__ = ['fit_spectra']

import sys
import copy
import pathlib
import inspect
from datetime import datetime
import multiprocessing as mp
import numpy as np
import eispac.core.fitting_functions as fit_fns
from eispac.core.eiscube import EISCube
from eispac.core.read_cube import read_cube
from eispac.core.read_template import EISFitTemplate
from eispac.core.eisfitresult import EISFitResult
from eispac.core.eisfitresult import create_fit_dict
from eispac.core.scale_guess import scale_guess
from eispac.core.mpfit import mpfit
from eispac.instr import calc_velocity

# i'm sure that there is a better way to do this!
cntr = mp.Value("i", 0) # a counter
nexp = mp.Value("i", 0) # total exposures

def check_for_name_guard(debug=False):
    """Check to see if the top level script has a "name guard" that will
    protect against multiprocessing spawning infinite child processes.
    """
    name_guard = False
    # Walk up the call stack and find the first program or script with
    # __name__ == "__main__". This should be the top-level program.
    if debug:
        print('')
        print('Checking call stack for a name guard...')
    call_stack = inspect.stack()
    for s in range(len(call_stack)):
        frame = call_stack[s][0]
        frame_filename = call_stack[s][1]
        if debug:
            print('Stack frame index:', s)
            print('   filename = ', frame_filename)
        if '__name__' in frame.f_locals:
            if debug:
                print('   __name__ = ', frame.f_locals['__name__'])
                print('   pathlib.Path.is_file() = ',
                      str(pathlib.Path(frame_filename).is_file()))
            if frame.f_locals['__name__'] == '__main__':
                if not frame_filename.endswith('.py'):
                    # Probably running in a Python terminal or interactive shell
                    name_guard = False
                else:
                    if pathlib.Path(frame_filename).is_file() == False:
                        # Probably running in an entry point executable
                        name_guard = True
                        break
                    # Examine the source code of __main__ script
                    with open(frame_filename, 'r') as f_file:
                        for line in f_file.readlines():
                            sl = line.replace(' ','') # Remove whitespace
                            sl = sl.replace('"', "'") # Convert " to '
                            if sl.startswith("if__name__=='__main__':"):
                                name_guard = True
                break

    return name_guard

def fit_with_mpfit(wave_cube, inten_cube, errs_cube, template, parinfo,
                   min_points=7,  chunk=1, data_units='unknown'):
    """Helper function for fit_spectra(). Fits one or more intensity spectra
    using the mpfit module.
    """

    with cntr.get_lock():
        cntr.value += 1
    #print(f' + working on exposure {cntr.value:03d}') #, end='\r')
    print(f' + working on exposure {chunk:03d}', end='\r')

    inten_size = inten_cube.shape
    n_pxls = inten_size[0]
    n_steps = inten_size[1]
    n_wave = inten_size[2]
    line_ids = template['line_ids']
    if isinstance(line_ids, bytes):
        line_ids = [line_ids.decode('utf-8')]

    # Extract fit information from template
    n_gauss = template['n_gauss']
    n_poly = template['n_poly']
    min_wave = template['data_x'][0]
    max_wave = template['data_x'][-1]
    oldguess = template['fit']

    # Create fit dictionary, mask array, and parameter indices
    fit_dict = create_fit_dict(n_pxls, n_steps, n_wave, n_gauss, n_poly,
                               data_units=data_units)
    fit_dict['line_ids'] = line_ids
    fit_dict['wave_range'][0] = min_wave
    fit_dict['wave_range'][1] = max_wave
    mask_ij = np.ones(n_wave)
    loc_peaks = np.arange(n_gauss)*3
    loc_cen = np.arange(n_gauss)*3+1
    loc_wid = np.arange(n_gauss)*3+2
    loc_backs = np.arange(n_poly)+n_gauss*3

    # loop over pixel positions and slit steps in the entire raster
    for ii in range(n_pxls):
        for jj in range(n_steps):

            # Extract a single profile from the raster
            wave_ij = wave_cube[ii,jj,::]
            inten_ij = inten_cube[ii,jj,::]
            errs_ij = errs_cube[ii,jj,::]

            # Cut data to only "good" nonzero errs within the wavelength range
            mask_ij[0:-1] = 1 # reset the mask (remember: 1 = True = masked val)
            loc_good = np.where((errs_ij > 0) & (wave_ij >= min_wave)
                                & (wave_ij <= max_wave))
            num_good_data = len(loc_good[0])
            if num_good_data > 0:
                mask_ij[loc_good] = 0 # unmask good data

            fit_dict['mask'][ii,jj,:] = mask_ij[:] # always save the mask

            if num_good_data < min_points:
                fit_dict['status'][ii,jj] = -1
                continue
            wave_ij = wave_ij[loc_good]
            inten_ij = inten_ij[loc_good]
            errs_ij = errs_ij[loc_good]

            # Scale guess parameters to data (should speed up fitting some)
            newguess = scale_guess(wave_ij, inten_ij, oldguess, n_gauss, n_poly)

            # Plug in new guess values to parinfo
            for i in range(len(newguess)):
                parinfo[i]['value'] = newguess[i]

            # Assemble dict of extra args to pass to mpfit
            fa = {'x': wave_ij, 'y': inten_ij, 'error': errs_ij,
                  'n_gauss': n_gauss, 'n_poly': n_poly}

            # Fit the profile
            out = mpfit(fit_fns.multigaussian_deviates, parinfo=parinfo,
                        functkw=fa, xtol=1.0E-6, ftol=1.0E-6, gtol=1.0E-6,
                        maxiter=2000, quiet=1)

            # check convergence status for a valid result
            if out.status > 0:
                # compute line inten and errors directly (may need to revisit)
                fpeaks = out.params[loc_peaks]
                fwdths = out.params[loc_wid]
                epeaks = out.perror[loc_peaks]
                ewdths = out.perror[loc_wid]
                l_inten = np.sqrt(2*np.pi)*fpeaks*fwdths
                e_inten = np.zeros(n_gauss)
                for n in range(n_gauss):
                    if fpeaks[n] != 0 and fwdths[n] != 0:
                        e_inten[n] = (l_inten[n]*np.sqrt((epeaks[n]/fpeaks[n])**2
                                                         +(ewdths[n]/fwdths[n])**2))
                    else:
                        e_inten[n] = 0.0

                # assemble fit structure
                fit_dict['status'][ii,jj] = out.status
                fit_dict['chi2'][ii,jj] = out.fnorm/out.dof
                fit_dict['wavelength'][ii,jj,:] = wave_cube[ii,jj,:]
                fit_dict['params'][ii,jj,:] = out.params
                fit_dict['perror'][ii,jj,:] = out.perror
                fit_dict['int'][ii,jj,:] = l_inten
                fit_dict['err_int'][ii,jj,:] = e_inten
            else:
                print(' ! fit did not converge!')
                fit_dict['status'][ii,jj] = out.status

    return fit_dict

def fit_spectra(inten, template, parinfo=None, wave=None, errs=None, min_points=7,
                ncpu='max', unsafe_mp=False, ignore_warnings=False,
                skip_fitting=False, debug=False):
    """Fit one or more EIS line spectra using mpfit (with multiprocessing).

    Parameters
    ----------
    inten : EISCube object, array_like, or filepath
        One or more intensity profiles to be fit. The code will loop over the data
        according to its dimensionality. 3D data is assumed to be a full EIS raster
        (or a sub region), 2D data is assumed to be a single EIS slit, and 1D data
        is assumed to be a single profile.
    template : EISFitTemplate object, dict, or filepath
        Either an EISFitTemplate, a 'template' dictionary, or the path to a
        template file.
    parinfo : list, optional
        List of dictionaries with fit parameters formatted for use with mpfit.
        Will supercede any parinfo lists loaded from an EISFitTemplate. Required
        if the 'template' parameter is given as a dictionary.
    wave : array_like, optional
        Associated wavelength values for the spectra. Required if 'inten' is
        given as an array and ignored otherwise.
    errs : array_like, optional
        Intensity error values for the spectra. Required if 'inten' is given as
        an array and ignored otherwise.
    min_points : int, optional
        Minimum number of good quality data points (i.e. non-zero values & errs)
        to be used in each fit. Spectra with fewer data points will be skipped.
        Must be a number >= the total number of fit parameters. Default is 7.
    ncpu : int, optional
        Number of cpu processes to parallelize over. Must be less than or equal
        to the total number of cores the system has. If set to 'max' or None, the
        code will use the maximum number of cores available. Default is 'max'.
        Important: due to the specifics of how the multiprocessing library works,
        any statements that call fit_spectra() using ncpu > 1 MUST be wrapped in
        a "if __name__ == __main__:" statement in the top-level program. If such
        a "name guard" statement is not detected, this function will fall back to
        using a single process.
    unsafe_mp : bool, optional
        If set to True, will use multiprocessing even if there is no name guard
        no name guard (if ncpu > 0). Used by the console script "eit_fit_files".
        Default is False (name guard enforced). Disabling the name guard runs the
        risk of spawning infinite processes if run incorrectly. USE AT YOUR OWN
        RISK!
    ignore_warnings : bool, optional
        If set to True, will silence the warning about a missing or disabled name
        guard (we are serious at it, be careful). Default is False.
    skip_fitting : bool, optional
        If set to True, will skip the fitting altogether and just return an empty
        EISFitResult instance. Used mainly for testing. Default is False.
    debug : bool, optional
        If set to True, will print some extra information useful for debugging
        development versions of the code. Default is False.

    Returns
    -------
    fit_res : EISFitResult class instance
        An EISFitResult object containing the output fit parameters.
    """
    # Validate template & parinfo and read / copy as needed
    if isinstance(template, (str, pathlib.Path)):
        template_obj = EISFitTemplate.read_template(template)
        if template_obj is None:
            return None
        template_copy = copy.deepcopy(template_obj.template)
        parinfo_copy = copy.deepcopy(template_obj.parinfo)
        tmplt_filename = template_obj.filename_temp
    elif isinstance(template, EISFitTemplate):
        template_copy = copy.deepcopy(template.template)
        parinfo_copy = copy.deepcopy(template.parinfo)
        tmplt_filename = template.filename_temp
    elif isinstance(template, dict):
        template_copy = copy.deepcopy(template)
        parinfo_copy = None
        tmplt_filename = 'unknown'
    else:
        print('Please input either the path to a template file, an'
             +' EISFitTemplate instance, or dictionary.', file=sys.stderr)
        return None

    if isinstance(parinfo, list):
        # Direct user-input always supercedes automatically loaded data
        parinfo_copy = copy.deepcopy(parinfo)
    elif parinfo_copy is None:
        print('Please input a parinfo list or a full EISFitTemplate.', file=sys.stderr)
        return None

    # If given an EISCube, extract data arrays and zero out masked values.
    # Otherwise, make local copies and just zero out negative values
    # TODO: Add more validation for the case of input arrays
    if isinstance(inten, (str, pathlib.Path)):
        central_wave = np.mean([template_copy['wmin'], template_copy['wmax']])
        eis_cube = read_cube(inten, window=float(central_wave))
        if eis_cube is None:
            return None
        wave_cube = eis_cube.wavelength.copy()
        errs_cube = eis_cube.uncertainty.array.copy()
        inten_cube = eis_cube.data.copy()
        metadata = copy.deepcopy(eis_cube.meta)
        data_units = eis_cube.unit.to_string()
        data_radcal = copy.deepcopy(eis_cube.radcal)
        loc_masked = np.where(eis_cube.mask == True)
        inten_cube[loc_masked] = 0
        errs_cube[loc_masked] = 0
        del eis_cube
    elif isinstance(inten, EISCube):
        wave_cube = inten.wavelength.copy()
        errs_cube = inten.uncertainty.array.copy()
        inten_cube = inten.data.copy()
        metadata = copy.deepcopy(inten.meta)
        data_units = inten.unit.to_string()
        data_radcal = copy.deepcopy(inten.radcal)
        loc_masked = np.where(inten.mask == True)
        inten_cube[loc_masked] = 0
        errs_cube[loc_masked] = 0
    elif isinstance(inten, np.ndarray):
        if not isinstance(wave, np.ndarray):
            print('Please input a wavelength array or a full EISCube.', file=sys.stderr)
            return None
        elif not isinstance(errs, np.ndarray):
            print('Please input an error array or a full EISCube.', file=sys.stderr)
            return None
        wave_cube = wave.copy()
        errs_cube = errs.copy()
        inten_cube = inten.copy()
        metadata = {'filename_data':'unknown', 'index':{}, 'pointing':{},
                    'radcal':'unknown', 'wave':'unknown'}
        data_units = 'unknown'
        data_radcal = 'unknown'
        loc_bad = np.where(errs_cube <= 0)
        inten_cube[loc_bad] = 0
        errs_cube[loc_bad] = 0
    else:
        print('Error: missing or invalid data. Please input a filepath, EISCube,'
             +' or complete set of intensity, wavelength, and error arrays.', file=sys.stderr)
        return None

    # Validate input ncpu value
    if str(ncpu).lower() == 'max' or str(ncpu).lower() == 'none':
        ncpu = mp.cpu_count()
    else:
        ncpu = int(ncpu)

    if ncpu <= 0:
        ncpu = 1
    elif ncpu > mp.cpu_count():
        ncpu = mp.cpu_count()

    # Ensure that multiprocessing will not spawn infinite child processes
    if ncpu > 1:
        name_guard = check_for_name_guard(debug=debug)
        if unsafe_mp == True and name_guard == False:
            if ignore_warnings == False:
                print('CRITICAL WARNING: unsafe_mp == True while no name guard'
                     +' was found in the top-level script! Be aware, parallel'
                     +' processes may freeze or behave unexpectedly.')
        elif name_guard == False:
            ncpu = 1
            if ignore_warnings == False:
                print('WARNING: no name guard was found in the top-level script!'
                     +' Falling back to a single process for safety.')

    # Check value of min_points
    n_params = len(parinfo_copy)
    if min_points is None or min_points < n_params:
        print('WARNING: min_points must be >= total number of fit parameters.'
             +' min_points has been set to '+str(n_params))
        min_points = n_params

    # Check the dimensions of input data.
    # If the the arrays are not 3D, add shallow dimensions of size 1
    num_dims = inten_cube.ndim
    dims_size = inten_cube.shape
    if num_dims == 1:
        n_pxls = 1
        n_steps = 1
        wave_cube = wave_cube[np.newaxis, np.newaxis, :]
        inten_cube = inten_cube[np.newaxis, np.newaxis, :]
        errs_cube = errs_cube[np.newaxis, np.newaxis, :]
    elif num_dims == 2:
        n_pxls = dims_size[0]
        n_steps = 1
        wave_cube = wave_cube[:, np.newaxis, :]
        inten_cube = inten_cube[:, np.newaxis, :]
        errs_cube = errs_cube[:, np.newaxis, :]
    elif num_dims == 3:
        n_pxls = dims_size[0]
        n_steps = dims_size[1]
        wave_cube = wave_cube
        inten_cube = inten_cube
        errs_cube = errs_cube

    # Initalize output object which will contain the fit results
    fit_res = EISFitResult(wave_cube, template_copy, parinfo_copy,
                           func_name='multigaussian', data_units=data_units,
                           radcal=data_radcal)
    fit_res.meta = metadata
    fit_res.meta['filename_template'] = tmplt_filename
    fit_res.fit_module = 'mpfit'
    fit_res.fit_method = 'LevMarLSQ'

    if skip_fitting != True:
        t1 = datetime.now() # start a simple timer
        nexp.value = n_steps
        print(f' + computing fits for {n_steps:d} exposures, each with {n_pxls:d} spectra')

        # Run fitting in either a single process (default) or using multiprocessing
        if ncpu == 1:
            print(' + running mpfit in a single process')
            fit_dict = fit_with_mpfit(wave_cube, inten_cube, errs_cube,
                                      template_copy, parinfo_copy, min_points,
                                      data_units=data_units)
            fit_res.fit = fit_dict
        else:
            if ncpu > n_steps:
                ncpu = n_steps
            print(f' + running mpfit on {ncpu:d} cores (of {mp.cpu_count():d})')

            # initialize pool of workers
            pool = mp.Pool(processes=ncpu)

            # Split out the data for each single slit and run the pool
            args = [(wave_cube[:,jj:jj+1,:], inten_cube[:,jj:jj+1,:], errs_cube[:,jj:jj+1,:],
                     template_copy, parinfo_copy, min_points, jj+1, data_units)
                     for jj in range(n_steps)]
            pool_out = pool.starmap(fit_with_mpfit, args)
            pool.close()

            # Now, loop over each slit and copy fit results values to full output object
            for jj in range(n_steps):
                fit_res.fit['status'][:,jj] = pool_out[jj]['status'][:,0]
                fit_res.fit['chi2'][:,jj] = pool_out[jj]['chi2'][:,0]
                fit_res.fit['mask'][:,jj,:] = pool_out[jj]['mask'][:,0,:]
                fit_res.fit['wavelength'][:,jj,:] = pool_out[jj]['wavelength'][:,0,:]
                fit_res.fit['params'][:,jj,:] = pool_out[jj]['params'][:,0,:]
                fit_res.fit['perror'][:,jj,:] = pool_out[jj]['perror'][:,0,:]
                fit_res.fit['int'][:,jj,:] = pool_out[jj]['int'][:,0,:]
                fit_res.fit['err_int'][:,jj,:] = pool_out[jj]['err_int'][:,0,:]

        # Calculate the Doppler velocity for each line
        # TODO: revisit error estimation
        for gg in range(fit_res.n_gauss):
            base_wave = parinfo_copy[1+3*gg]['value']
            obs_cent = fit_res.fit['params'][:,:,1+3*gg]
            obs_errs = fit_res.fit['perror'][:,:,1+3*gg]
            velocity = calc_velocity(obs_cent, base_wave)
            fit_res.fit['vel'][:,:,gg] = velocity
            rel_err = obs_errs/obs_cent
            fit_res.fit['err_vel'][:,:,gg] = rel_err*velocity

        # print status
        t2 = datetime.now() # end timer
        num_fit = len(np.where(fit_res.fit['status'] > -0)[0])
        num_too_few = len(np.where((fit_res.fit['status'] == -2) |
                                   (fit_res.fit['status'] == -1))[0])
        num_bad_params = len(np.where((fit_res.fit['status'] == -3) |
                                      (fit_res.fit['status'] == 0))[0])
        print('\n')
        print('Finished computing fits!')
        print(f'   runtime : {t2-t1}')
        print(f'   {num_fit} spectra fit without issues')
        print(f'   {num_too_few} spectra have < {min_points} good data points')
        print(f'   {num_bad_params} spectra have bad or invalid parameters')

        # reset global counters
        cntr.value = 0
        nexp.value = 0

    return fit_res

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import astropy.units as u
    # from read_cube import read_cube
    # from read_template import read_template

    # input data and template files
    # file_data = './data/eis_20120924_105026.head.h5'
    file_data = './data/eis_20190404_131513.data.h5'
    file_template = './templates/eis_template_dir/fe_12_195_119.2c.template.h5'

    # read fit template
    Fe_XII_195_119 = EISFitTemplate.read_template(file_template)

    # read spectra window
    raster = read_cube(file_data, Fe_XII_195_119.central_wave)

    # fit profile
    # wave_coords = raster.axis_world_coords('em.wl')
    # lower_corner = (300*u.arcsec, 50*u.arcsec, wave_coords[0])
    # upper_corner = (400*u.arcsec, 150*u.arcsec, wave_coords[-1])
    # sub_raster = raster.crop_by_coords(lower_corner, upper_corner=upper_corner)
    sub_raster = raster[0:10, 0:10, :]

    fit_res = fit_spectra(sub_raster, Fe_XII_195_119, ncpu=4)

    # Quick plot raster
    plot_aspect_ratio = raster.meta['pointing']['y_scale']/raster.meta['pointing']['x_scale']
    raster.sum_spectra().plot(aspect=plot_aspect_ratio)

    # Plot example fit
    ex_pxl_coords = [5, 5]
    # ex_pxl_coords = [4, 8]
    fit_x, fit_y = fit_res.get_fit_profile(coords=ex_pxl_coords, num_wavelengths=100)
    c0_fit_x, c0_fit_y = fit_res.get_fit_profile(component=0, coords=ex_pxl_coords,
                                                 num_wavelengths=100)
    c1_fit_x, c1_fit_y = fit_res.get_fit_profile(component=1, coords=ex_pxl_coords,
                                                 num_wavelengths=100)
    c2_fit_x, c2_fit_y = fit_res.get_fit_profile(component=2, coords=ex_pxl_coords,
                                                 num_wavelengths=100)
    sub_data = sub_raster.data[ex_pxl_coords[0], ex_pxl_coords[1], :]
    sub_wave = sub_raster.wavelength[ex_pxl_coords[0], ex_pxl_coords[1], :]
    sub_err = sub_raster.uncertainty.array[ex_pxl_coords[0], ex_pxl_coords[1], :]

    fig = plt.figure()
    profile_subplt = fig.add_subplot(111)
    profile_subplt.errorbar(sub_wave, sub_data, yerr=sub_err,
                            ls='', marker='o', color='k')
    profile_subplt.plot(fit_x, fit_y, color='b')
    profile_subplt.plot(c0_fit_x, c0_fit_y, color='r')
    profile_subplt.plot(c1_fit_x, c1_fit_y, color='r', ls='--')
    profile_subplt.plot(c2_fit_x, c2_fit_y, color='g')
    profile_subplt.set_xlabel(r'Wavelength [$\AA$]')
    profile_subplt.set_ylabel('Intensity ['+raster.unit.to_string()+']')
    plt.show()


##Remaining tasks: -ve value of velocity, maxlimit of velocity, 0-acceleration
import numpy as np
import math
import copy
import dubins
import shapely.geometry as geom
import threading
from statistics import median

#Change radius of curvature for 0.9
from vel_acc_to_throttle import *

lock = threading.Lock()
inf = 1e9
No_of_threads = 11

acc= {}
acc[0] = [-1.0, -0.5, 0.5, 1.0, 2.0, 4.0]
acc[1] = [-1.0, -0.5, 0.5, 1.0, 2.0, 4.0]
acc[2] = [-1.0, -0.5, 0.5, 1.0, 2.0, 4.0]
acc[3] = [-1.0, -0.5, 0.5, 1.0, 2.0, 4.0]
acc[4] = [-1.0, -0.5, 0.5, 1.0, 2.0, 4.0]
acc[5] = [-1.0, 0.0, 1.0, 2.0, 4.0]
acc[6] = [-1.0, 0.0, 1.0, 2.0, 4.0]
acc[7] = [-1.0, 0.0, 1.0, 2.0, 4.0]
acc[8] = [-1.0, 0.0, 1.0, 2.0, 4.0]
acc[9] =  [-3.0, -1.0, 0.0, 1.0, 2.0, 4.0]
acc[10] = [-3.0, -1.0, 0.0, 1.0, 2.0, 4.0]
acc[11] = [-3.0, -1.0, 0.0, 1.0, 2.0, 4.0]
acc[12] = [-3.0, -1.0, 0.0, 1.0, 2.0, 4.0]
acc[13] = [-3.0, -1.0, 0.0, 1.0, 2.0, 4.0]
acc[14] = [-3.0, -1.0, 0.0, 1.0, 2.0, 4.0]
acc[15] = [-3.0, -1.0, 0.0, 1.0, 2.0, 4.0]
acc[16] = [-5.0, -3.0, -1.0, 0.0, 1.5, 3.0]
acc[17] = [-5.0, -3.0, -1.0, 0.0, 1.5, 3.0]
acc[18] = [-5.0, -3.0, -1.0, 0.0, 1.5, 3.0]
acc[19] = [-5.0, -3.0, -1.0, 0.0, 1.5, 3.0]
acc[20] = [-5.0, -3.0, -1.0, 0.0, 1.5, 3.0]
acc[21] = [-5.0, -3.0, -1.0, 0.0, 1.5]
acc[22] = [-5.0, -3.0, -1.2, 0.0, 1.5]
acc[23] = [-5.0, -3.0, -1.2, 0.0, 1.0]
acc[24] = [-5.0, -3.0, -1.2, 0.0, 1.0]
acc[25] = [-5.0, -3.0, -1.2, 0.0, 1.0]
acc[26] = [-5.0, -3.0, -1.3, 0.0, 1.0]
acc[27] = [-5.0, -3.0, -1.4, 0.0, 1.0]
acc[28] = [-5.0, -3.0, -1.4, 0.0, 0.5]
acc[29] = [-5.0, -3.0, -1.4, 0.0, 0.5]
acc[30] = [-5.0, -3.0, -1.5, 0.0]

total_distance = 150.0
grid_points = []

actual_vel = {} #key = (i,j,v,t)
actual_tim = {} #key = (i,j,v,t)
prev_acc = {} #key = (i,j,v,t)
c = {} #key = (j,v,t)
p = {} #key = (i,j,v,t)
velocities = []
times= []

#Used for updation across different layers
temp_tim = {}
temp_vel = {}
temp_acc = {}
temp_c = {}
temp_p = {}

temp_theta = {}
cur_theta = {}
#-----------------------------------------

y_step = 0.6
x_step = 5.0
w = 3.6
obs_initial_pos = [450.0,0.0]
obs_vel = 5.0

corner_local_coords = [[-2.5, -1.1], [-2.5, 1.1], [2.5, 1.1], [2.5, -1.1]]

Radius_of_road = 20.0

Obs2_initial_pos = [360.0, 2.4]
Obs2_vel = 5.0

def RadiusofCurvature(start_pt, end_pt, turn_radius=30.0, step_size=1.0):
    """Generate points along a Dubins path connecting start point to end point.
    Format for input / output points: (x, y, angle)"""
    min_turn_radius = min(0.1, turn_radius)
    satisfied = False
    configurations = [start_pt, end_pt]
    while not satisfied:
        dubins_path = dubins.shortest_path(start_pt, end_pt, turn_radius)
        configurations, _ = dubins_path.sample_many(step_size)
        cex_found = False
        for configuration in configurations:
            if not (min(start_pt[0], end_pt[0]) - 0.1 <= configuration[0] <= max(start_pt[0], end_pt[0]) + 0.1 and
                    min(start_pt[1], end_pt[1]) - 0.1 <= configuration[1] <= max(start_pt[1], end_pt[1]) + 0.1):
                cex_found = True
                break
        satisfied = not cex_found
        if cex_found:
            # Decrease radius until finding a satisfying result.
            # We could do a binary search but that requires a termination condition.
            turn_radius = turn_radius*0.8
            if turn_radius < min_turn_radius:
                break
    if not satisfied:
        return 0.1
    return turn_radius

def rotate_point_cw(point, theta):
    cos_theta = math.cos(theta)
    sin_theta = math.sin(theta)
    return np.dot(np.array([[cos_theta, sin_theta], [-sin_theta, cos_theta]]), point)

def ObsPosition(t):
    Total_time = (1000.0 + 2*math.pi*20.0)/obs_vel
    t = t - Total_time * int(t/Total_time)
    offset = t * obs_vel
    if( obs_initial_pos[0] - offset >=0):
        return [obs_initial_pos[0]-offset,obs_initial_pos[1], math.pi]
    elif( obs_initial_pos[0] - (offset - math.pi * Radius_of_road) >=0 ):
        turned_theta = (offset - obs_initial_pos[0])/Radius_of_road
        return [-Radius_of_road*math.sin(turned_theta), -Radius_of_road + Radius_of_road*math.cos(turned_theta), math.pi + turned_theta]
    elif( offset <= obs_initial_pos[0] + 500.0 + Radius_of_road*math.pi):
        return [offset - obs_initial_pos[0] - Radius_of_road*math.pi, -2*Radius_of_road, 0.0]
    elif(offset <= 2*Radius_of_road*math.pi + obs_initial_pos[0]+500.0):
        turned_theta = (offset - Radius_of_road*math.pi - 500.0 - obs_initial_pos[0])/Radius_of_road
        return [500.0+Radius_of_road*math.sin(turned_theta),-Radius_of_road- Radius_of_road*math.cos(turned_theta), turned_theta]
    else:
        return [1000.0 - offset + 2*Radius_of_road*math.pi + obs_initial_pos[0], 0.0, math.pi]

def computeD(x1,x2,y1,y2,xp,yp):
    D = (x2 - x1) * (yp - y1) - (xp - x1) * (y2 - y1)
    return D

def check_colliding(pt2):
    turned = pt2[0][2]-pt2[4]
    # obstacle_position = [obs_initial_pos[0] - obs_vel*pt2[3],obs_initial_pos[1]]
    obstacle_position =  ObsPosition(pt2[3])
    car_corner_pos = []
    for local_coord in corner_local_coords:
        rotated_local_coord = \
            rotate_point_cw(point=np.transpose(np.array(local_coord)),
                             theta=pt2[4])

        car_corner_pos.append([pt2[0][0]+rotated_local_coord[0],pt2[0][1]+rotated_local_coord[1]])

    # print(car_corner_pos)

    obs_corner_pos = []
    for local_coord in corner_local_coords:
        # rotated_local_coord = \
        #     rotate_point_ccw(point=np.transpose(np.array(local_coord)),
        #                      rotation_angle=-detected_objects[obj_ind].object_yaw_angle)
        rotated_local_coord = \
            rotate_point_cw(point=np.transpose(np.array(local_coord)),
                             theta=pt2[0][2])

        obs_corner_pos.append([obstacle_position[0] + rotated_local_coord[0],
                             obstacle_position[1] + rotated_local_coord[1]])

        # print(rotated_local_coord)
    # print(obs_corner_pos)

    collision = 0
    for dx in np.arange(-max(pt2[2],10),max(pt2[2],10),4.9):
        for pos in car_corner_pos:
            x = pos[0] + dx*math.cos(pt2[4])
            y = pos[1] + dx*math.sin(pt2[4])
            D1 = computeD(obs_corner_pos[0][0],obs_corner_pos[1][0],obs_corner_pos[0][1],obs_corner_pos[1][1],x,y)
            D2 = computeD(obs_corner_pos[2][0],obs_corner_pos[3][0],obs_corner_pos[2][1],obs_corner_pos[3][1],x,y)
            D3 = computeD(obs_corner_pos[1][0],obs_corner_pos[2][0],obs_corner_pos[1][1],obs_corner_pos[2][1],x,y)
            D4 = computeD(obs_corner_pos[3][0],obs_corner_pos[0][0],obs_corner_pos[3][1],obs_corner_pos[0][1],x,y)
            if ( D1*D2 >=0 and D3*D4>=0):
                collision=1
                break

    if(collision == 1):
        return collision

    obstacle_position = [Obs2_initial_pos[0] + Obs2_vel*pt2[3],Obs2_initial_pos[1]]
    obs_corner_pos = []
    for local_coord in corner_local_coords:
        # rotated_local_coord = \
        #     rotate_point_ccw(point=np.transpose(np.array(local_coord)),
        #                      rotation_angle=-detected_objects[obj_ind].object_yaw_angle)
        rotated_local_coord = \
            rotate_point_cw(point=np.transpose(np.array(local_coord)),
                             theta=pt2[0][2])

        obs_corner_pos.append([obstacle_position[0] + rotated_local_coord[0],
                             obstacle_position[1] + rotated_local_coord[1]])

    for dx in np.arange(-max(pt2[2],10),max(pt2[2],10),4.9):
        for pos in car_corner_pos:
            x = pos[0] + dx*math.cos(pt2[4])
            y = pos[1] + dx*math.sin(pt2[4])
            D1 = computeD(obs_corner_pos[0][0],obs_corner_pos[1][0],obs_corner_pos[0][1],obs_corner_pos[1][1],x,y)
            D2 = computeD(obs_corner_pos[2][0],obs_corner_pos[3][0],obs_corner_pos[2][1],obs_corner_pos[3][1],x,y)
            D3 = computeD(obs_corner_pos[1][0],obs_corner_pos[2][0],obs_corner_pos[1][1],obs_corner_pos[2][1],x,y)
            D4 = computeD(obs_corner_pos[3][0],obs_corner_pos[0][0],obs_corner_pos[3][1],obs_corner_pos[0][1],x,y)
            if ( D1*D2 >=0 and D3*D4>=0):
                collision=1
                break

    return collision

def cost(c1, pt1,pt2, off=0.0):
    # print(pt1)
    # print(pt2)
    # r = RadiusofCurvature(pt1[0],pt2[0])
    R={}
    R[(5,0)] = inf
    # For straight line only

    deltay = abs(pt2[0][1]-pt1[0][1])
    deltax = abs(pt2[0][0]-pt1[0][0])
    temp = (deltax,deltay)
    if(temp in R):
        r = R[temp]
    else:
        r = RadiusofCurvature([pt1[0][0],pt1[0][1],pt1[4]],[pt2[0][0],pt2[0][1],pt2[4]])
        if(r==30):
            r=inf
        R[temp] = r

    obstacle_position = [obs_initial_pos[0] - obs_vel*pt2[3],obs_initial_pos[1]]

    static_cost =  c1 + math.sqrt((pt2[0][0]-pt1[0][0])**2 + (pt2[0][1]-pt1[0][1])**2) + 10.0/r + 1.0*abs(off) + 0.1*math.exp(-0.1*math.sqrt((pt2[0][0]-obstacle_position[0])**2 + (pt2[0][1]-obstacle_position[1])**2))

    dynamic_cost = 15.0*(pt2[3]-pt1[3]) + (pt2[2]**2)*0.0 + 0.0*(pt2[1]**2) + 1.7e-10*(((pt2[1]-pt1[1])/(pt2[3]-pt1[3]))**2) + 1.0*(((pt2[2])**2)/r)

    return static_cost + dynamic_cost + check_colliding(pt2)*inf

    #off = 1 or 0.5
def Grid1(cur_pt,dist_to_cover):
    global grid_points
    x1 = round(cur_pt[0],2)
    x2 = max(x1-dist_to_cover,0) ##path to travel in first part of the road
    for i in np.arange(x1,x2,-x_step):
        gp = []
        for j in np.arange(w,-w,-y_step):
            gp.append([i,round(j,2),math.pi])
        grid_points.append(gp)
    return dist_to_cover - (x1-x2)

def Grid2(cur_pt,dist_to_cover):
    global grid_points
    theta_covered = math.atan(abs(cur_pt[0])/(Radius_of_road+cur_pt[1]))
    if(theta_covered<0):
        theta_covered = theta_covered + math.pi
    theta_to_cover = dist_to_cover/Radius_of_road
    final_theta = min(theta_covered + theta_to_cover,math.pi)
    for theta in np.arange(theta_covered,final_theta+0.00001,x_step/Radius_of_road):
        gp = []
        for j in np.arange(Radius_of_road+w,Radius_of_road-w,-y_step):
            x_coord = round(-j*math.sin(theta),2)
            y_coord = round(-Radius_of_road+j*math.cos(theta),2)
            gp.append([x_coord,y_coord,math.pi+theta])
        grid_points.append(gp)
    return (theta_covered + theta_to_cover - final_theta)*Radius_of_road

def Grid3(cur_pt,dist_to_cover):
    global grid_points
    x1 = round(cur_pt[0],2)
    x2 = min(x1+dist_to_cover,500.0) ##path to travel in first part of the road
    for i in np.arange(x1,x2,x_step):
        gp = []
        for j in np.arange(-2*Radius_of_road + w,-2*Radius_of_road-w,-y_step):
            gp.append([i,round(j,2),0.0])
        grid_points.append(gp)
    return (dist_to_cover - (x2-x1))

def Grid4(cur_pt,dist_to_cover):
    global grid_points
    theta_covered = math.atan(abs(cur_pt[0]-500.0)/(-Radius_of_road-cur_pt[1]))
    if(theta_covered<0):
        theta_covered = theta_covered + math.pi
    theta_to_cover = dist_to_cover/Radius_of_road
    final_theta = min(theta_covered + theta_to_cover,math.pi)
    for theta in np.arange(theta_covered,final_theta+0.0000001,x_step/Radius_of_road):
        gp = []
        for j in np.arange(Radius_of_road+w,Radius_of_road-w,-y_step):
            x_coord = round(500.0+j*math.sin(theta),2)
            y_coord = round(-Radius_of_road-j*math.cos(theta),2)
            gp.append([x_coord,y_coord,theta])
        grid_points.append(gp)
    return (theta_covered + theta_to_cover - final_theta)*Radius_of_road

def calculate_grid(cur_pt,dist_to_cover):
    global grid_points
    grid_points = []
    if(cur_pt[0]>0 and cur_pt[0]<=500 and cur_pt[1]>-20.0):  ##check in first part of the road
        remaining_dist = Grid1(cur_pt,dist_to_cover)
        if(remaining_dist > 0):
            remaining_dist = Grid2([0.0,0.0],remaining_dist)
        if(remaining_dist > 0):
            remaining_dist = Grid3([0.0,-2*Radius_of_road],remaining_dist)
        if(remaining_dist > 0):
            remaining_dist = Grid4([500.0,-2*Radius_of_road],remaining_dist)
    elif(cur_pt[0]<=0):
        remaining_dist = Grid2(cur_pt,dist_to_cover)
        if(remaining_dist>0):
            remaining_dist = Grid3([0.0,-2*Radius_of_road],remaining_dist)
        if(remaining_dist > 0):
            remaining_dist = Grid4([500.0,-2*Radius_of_road],remaining_dist)
    elif(cur_pt[0]>=0 and cur_pt[0]<500 and cur_pt[1]<-20.0):
        remaining_dist = Grid3(cur_pt,dist_to_cover)
        if(remaining_dist > 0):
            remaining_dist = Grid4([500.0,-2*Radius_of_road],remaining_dist)
        if(remaining_dist > 0):
            remaining_dist = Grid1([500.0,0.0],remaining_dist)
    else:
        remaining_dist = Grid4([500.0,-2*Radius_of_road],dist_to_cover)
        if(remaining_dist > 0):
            remaining_dist = Grid1([500.0,0.0],remaining_dist)

def computeTargetPath(cur_pt, dist_to_cover):

    calculate_grid(cur_pt,dist_to_cover)
    global grid_points
    global c
    global p
    global actual_vel
    global actual_tim
    global prev_acc
    global cur_theta

    # print(grid_points)

    ##########change from here
    X = round(2*w/y_step)
    Y = len(grid_points)

    ind2 = -1
    min_dist = inf
    for j in range(X):
        cur_dist = (grid_points[0][j][1]-cur_pt[1])**2 + (grid_points[0][j][0]-cur_pt[0])**2
        if(cur_dist < min_dist):
            min_dist = cur_dist
            ind2 = j

    #Initialisation
    i3 = math.ceil(cur_pt[3])
    i4 = math.ceil(cur_pt[4])
    c[(ind2,i3,i4)] = 0.0
    p[(0,ind2,i3,i4)] = -1
    actual_tim[(0,ind2,i3,i4)] = cur_pt[4]
    actual_vel[(0,ind2,i3,i4)] = cur_pt[3]
    prev_acc[(0,ind2,i3,i4)] = cur_pt[2]
    cur_theta[(0,ind2,i3,i4)] = cur_pt[5]

    global velocities
    global times
    global temp_vel
    global temp_c
    global temp_tim
    global temp_p
    global temp_acc
    global temp_theta

    cf = inf
    final_pos = -1

    for i in  range(Y-1):
        t0= threading.Thread(target=parallel_func, args=(0,i,X,))
        t1= threading.Thread(target=parallel_func, args=(1,i,X,))
        t2= threading.Thread(target=parallel_func, args=(2,i,X,))
        t3= threading.Thread(target=parallel_func, args=(3,i,X,))
        t4= threading.Thread(target=parallel_func, args=(4,i,X,))
        t5= threading.Thread(target=parallel_func, args=(5,i,X,))
        t6= threading.Thread(target=parallel_func, args=(6,i,X,))
        t7= threading.Thread(target=parallel_func, args=(7,i,X,))
        t8= threading.Thread(target=parallel_func, args=(8,i,X,))
        t9= threading.Thread(target=parallel_func, args=(9,i,X,))
        t10= threading.Thread(target=parallel_func, args=(10,i,X,))
        t0.start()
        t1.start()
        t2.start()
        t3.start()
        t4.start()
        t5.start()
        t6.start()
        t7.start()
        t8.start()
        t9.start()
        t10.start()
        t0.join()
        t1.join()
        t2.join()
        t3.join()
        t4.join()
        t5.join()
        t6.join()
        t7.join()
        t8.join()
        t9.join()
        t10.join()

        # print(velocities)
        # print(" ")
        v_m = median(velocities)
        t_m = median(times)
        v_min = v_m-5
        v_max = v_m+5
        t_max = t_m+5
        t_min = t_m-5

        # print(c)
        c = {}
        for (j,v,t) in temp_c:
            ind_v = math.ceil(v)
            if(v > v_max):
                ind_v = inf
            if(v < v_min):
                ind_v = v_min
            ind_t = math.ceil(t)
            if(t > t_max):
                ind_t = inf
            if(t < t_min):
                ind_t = t_min

            if ((j,ind_v,ind_t) not in c) or (c[(j,ind_v,ind_t)] > temp_c[(j,v,t)] ):
                c[(j,ind_v,ind_t)] = temp_c[(j,v,t)]
                p[(i+1,j,ind_v,ind_t)] = temp_p[(i+1,j,v,t)]
                actual_vel[(i+1,j,ind_v,ind_t)] = temp_vel[(i+1,j,v,t)]
                actual_tim[(i+1,j,ind_v,ind_t)] = temp_tim[(i+1,j,v,t)]
                prev_acc[(i+1,j,ind_v,ind_t)] = temp_acc[(i+1,j,v,t)]
                cur_theta[(i+1,j,ind_v,ind_t)] = temp_theta[(i+1,j,v,t)]
                if(i==Y-2) and (cf>c[(j,ind_v,ind_t)]):
                    cf = c[(j,ind_v,ind_t)]
                    final_pos = (i+1,j,ind_v,ind_t)

        velocities = []
        times = []
        temp_c = {}
        temp_vel = {}
        temp_acc = {}
        temp_p = {}
        temp_tim = {}
        temp_theta = {}

    travel_path = []
    (i,j,ind2,ind3) = final_pos
    while ( (p[(i,j,ind2,ind3)]) != -1 ):
        travel_path = [[float(grid_points[i][j][0]),float(grid_points[i][j][1]),prev_acc[(i,j,ind2,ind3)],actual_vel[(i,j,ind2,ind3)],actual_tim[(i,j,ind2,ind3)],cur_theta[(i,j,ind2,ind3)]]] + travel_path
        (i,j,ind2,ind3) = (p[(i,j,ind2,ind3)])

    return travel_path

def parallel_func(ind4,i,X):
    global c
    global p
    global actual_vel
    global actual_tim
    global prev_acc

    global temp_c
    global temp_p
    global temp_acc
    global temp_vel
    global temp_tim
    global temp_theta

    global velocities
    global times
    global lock

    for (j,ind2,ind3) in c:

        v_i = math.ceil(actual_vel[(i,j,ind2,ind3)])
        if(ind4 < len(acc[v_i])):
            m1 = max(0,j-1)
            m2 = min(9,j+1)
            for k in range(m1,m2+1):
                a_f = acc[v_i][ind4]
                cur_cost = 0
                v_f = ( (actual_vel[(i,j,ind2,ind3)]**2) +2*a_f*x_step)
                if(v_f < 0):
                    continue
                else:
                    v_f = v_f ** 0.5
                if(v_f > 30):
                    continue
                v_f = round(v_f,4)

                ind5 = math.ceil(v_f)
                if v_f == actual_vel[(i,j,ind2,ind3)]:
                    t_f = x_step/v_f + actual_tim[(i,j,ind2,ind3)]
                else:
                    t_f = (v_f-actual_vel[(i,j,ind2,ind3)])/a_f + actual_tim[(i,j,ind2,ind3)]
                t_f = round(t_f,2)
                ind6 = math.ceil(t_f)

                x1 = grid_points[i][j][0]
                y1 = grid_points[i][j][1]
                x2 = grid_points[i+1][k][0]
                y2 = grid_points[i+1][k][1]
                if(x2 < x1):
                    if(y2 >= y1):
                        curtheta = math.pi + math.atan((y2-y1)/(x2-x1))
                    else:
                        curtheta = math.pi + math.atan((y2-y1)/(x2-x1))
                elif(x2 > x1):
                    if(y2 >= y1):
                        curtheta = math.atan((y2-y1)/(x2-x1))
                    else:
                        curtheta = 2*math.pi + math.atan((y2-y1)/(x2-x1))
                else:
                    if(y2>y1):
                        curtheta = math.pi/2.0
                    else:
                        curtheta = 1.5*math.pi

                # curtheta = grid_points[i+1][k][2] - math.atan((k-j)*y_step/x_step)

                cur_cost = cost(c[(j,ind2,ind3)],(grid_points[i][j],prev_acc[(i,j,ind2,ind3)],actual_vel[(i,j,ind2,ind3)],actual_tim[(i,j,ind2,ind3)], cur_theta[(i,j,ind2,ind3)]),(grid_points[i+1][k],a_f,v_f,t_f,curtheta),off=abs(w-k*y_step))
                if(cur_cost > inf):
                    continue
                velocities.append(v_f)
                times.append(t_f)
                lock.acquire(True)
                if( (k,ind5,ind6) not in temp_c) or (temp_c[(k,ind5,ind6)] > cur_cost):
                    temp_tim[(i+1,k,ind5,ind6)] = t_f
                    temp_c[(k,ind5,ind6)] = cur_cost
                    temp_vel[(i+1,k,ind5,ind6)] = v_f
                    temp_acc[(i+1,k,ind5,ind6)] = a_f
                    temp_p[(i+1,k,ind5,ind6)] = (i,j,ind2,ind3)
                    temp_theta[(i+1,k,ind5,ind6)] = curtheta
                lock.release()

# cur_pt = [16.77,0.0,0.5,34.45,26.0, math.pi]
# cur_pt =  [500.0, 0.0, 0.0, 0.0, 0.0, math.pi]
# cur_pt = [[405.0, 0.0, math.pi], 1.5, 16.583, 8.9, math.pi]
# c = check_colliding(cur_pt)
# print(c)

cur_pt =  [500.0, 0.0, 0.0, 10.0, 0.0, math.pi]
path = [cur_pt]
total_distance_covered = 0
while(total_distance_covered < 400):
    # path = path + computeTargetPath(cur_pt,100)
    future_path = computeTargetPath(cur_pt,100)
    path = path + future_path[:10]
    print(path)
    # print("path=====================")
    # print(path)
    total_distance_covered = 50.0 + total_distance_covered
    cur_pt = path[-1]
    actual_vel = {}
    actual_tim = {}
    prev_acc = {}
    c = {}
    p = {}
    # print(cur_pt)
    # print(path)

output = path
print(output)
print(" ")
target_path = []
# v = []
t = []
# a= []
throttle = []
prev = -1

for i in output:
    target_path.append([i[0],i[1]])
    # a.append(i[2])
    # v.append(i[3])

    if(prev == -1):
        prev = (i[3],i[2])
    else:
        t.append(i[4])
        throttle.append(throttle_value( (i[3]+prev[0])/2.0,i[2]))
        prev = (i[3],i[2])

print(throttle)
print(" ")
print(target_path)
print(" ")
print(t)

# r= RadiusofCurvature([5.0,0.0,math.pi],[0.0,0.25,math.pi])
# print(r)


# Copyright Contributors to the Pyro project.
# SPDX-License-Identifier: Apache-2.0

from abc import ABCMeta
from collections import namedtuple
import inspect

from jax import random, tree_map, vmap
from jax.flatten_util import ravel_pytree
import jax.numpy as jnp
import tensorflow_probability.substrates.jax as tfp

from numpyro.infer import init_to_uniform
from numpyro.infer.mcmc import MCMCKernel
from numpyro.infer.util import initialize_model
from numpyro.util import identity

TFPKernelState = namedtuple('TFPKernelState', ['z', 'kernel_results', 'rng_key'])

def _extract_kernel_functions(kernel):

    def init_fn(z, rng_key):
        z_flat, _ = ravel_pytree(z)
        results = kernel.bootstrap_results(z_flat)
        return TFPKernelState(z, results, rng_key)

    def sample_fn(state, model_args=(), model_kwargs=None):
        rng_key, rng_key_transition = random.split(state.rng_key)
        z_flat, unravel_fn = ravel_pytree(state.z)
        z_new_flat, results = kernel.one_step(z_flat, state.kernel_results, seed=rng_key_transition)
        return TFPKernelState(unravel_fn(z_new_flat), results, rng_key)

    return init_fn, sample_fn

def _make_log_prob_fn(potential_fn, unravel_fn):

    def log_prob_fn(x):
        # we deal with batched x in case the kernel is ReplicaExchangeMC
        batch_shape = jnp.shape(x)[:-1]
        if batch_shape:
            flatten_result = vmap(lambda a: -potential_fn(unravel_fn(a)))(
                jnp.reshape(x, (-1,) + jnp.shape(x)[-1:]))
            return tree_map(lambda a: jnp.reshape(a, batch_shape + jnp.shape(a)[1:]),
                            flatten_result)
        else:
            return - potential_fn(unravel_fn(x))

    return log_prob_fn

class _TFPKernelMeta(ABCMeta):
    def __getitem__(cls, kernel_class):
        assert issubclass(kernel_class, tfp.mcmc.TransitionKernel)
        assert 'target_log_prob_fn' in inspect.getfullargspec(kernel_class).args, \
            f"the first argument of {kernel_class} must be `target_log_prob_fn`"

        _PyroKernel = type(kernel_class.__name__, (TFPKernel,), {})
        _PyroKernel.kernel_class = kernel_class
        return _PyroKernel

class TFPKernel(MCMCKernel, metaclass=_TFPKernelMeta):
    """
    A thin wrapper for TensorFlow Probability (TFP) MCMC transition kernels.
    The argument `target_log_prob_fn` in TFP is replaced by either `model`
    or `potential_fn` (which is the negative of `target_log_prob_fn`).

    This class can be used to convert a TFP kernel to a NumPyro-compatible one
    as follows::

        kernel = TFPKernel[tfp.mcmc.NoUTurnSampler](model, step_size=1.)

    .. note:: By default, uncalibrated kernels will be inner kernels of the
        :class:`~tensorflow_probability.substrates.jax.mcmc.MetropolisHastings` kernel.

    .. note:: For :class:`~numpyro.contrib.tfp.mcmc.ReplicaExchangeMC`, TFP requires
        that the shape of `step_size` of the inner kernel must be
        `[len(inverse_temperatures), 1]` or `[len(inverse_temperatures), latent_size]`.

    :param model: Python callable containing Pyro :mod:`~numpyro.primitives`.
        If model is provided, `potential_fn` will be inferred using the model.
    :param potential_fn: Python callable that computes the target potential energy
        given input parameters. The input parameters to `potential_fn`
        can be any python collection type, provided that `init_params` argument to
        :meth:`init` has the same type.
    :param callable init_strategy: a per-site initialization function.
        See :ref:`init_strategy` section for available functions.
    :param kernel_kwargs: other arguments to be passed to TFP kernel constructor.
    """
    kernel_class = None

    def __init__(self, model=None, potential_fn=None, init_strategy=init_to_uniform,
                 **kernel_kwargs):
        if not (model is None) ^ (potential_fn is None):
            raise ValueError('Only one of `model` or `potential_fn` must be specified.')
        self._model = model
        self._potential_fn = potential_fn
        self._kernel_kwargs = kernel_kwargs
        self._init_strategy = init_strategy
        # Set on first call to init
        self._init_fn = None
        self._postprocess_fn = None
        self._sample_fn = None

    def _init_state(self, rng_key, model_args, model_kwargs, init_params):
        if self._model is not None:
            init_params, potential_fn, postprocess_fn, model_trace = initialize_model(
                rng_key,
                self._model,
                init_strategy=self._init_strategy,
                dynamic_args=True,
                model_args=model_args,
                model_kwargs=model_kwargs)
            init_params = init_params.z
            if self._init_fn is None:
                _, unravel_fn = ravel_pytree(init_params)
                kernel = self.kernel_class(
                    _make_log_prob_fn(potential_fn(*model_args, **model_kwargs), unravel_fn),
                    **self._kernel_kwargs)
                # Uncalibrated... kernels have to used inside MetropolisHastings, see
                # https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/mcmc/UncalibratedLangevin
                if self.kernel_class.__name__.startswith("Uncalibrated"):
                    kernel = tfp.mcmc.MetropolisHastings(kernel)
                self._init_fn, self._sample_fn = _extract_kernel_functions(kernel)
            self._postprocess_fn = postprocess_fn
        elif self._init_fn is None:
            _, unravel_fn = ravel_pytree(init_params)
            kernel = self.kernel_class(
                _make_log_prob_fn(self._potential_fn, unravel_fn),
                **self._kernel_kwargs)
            if self.kernel_class.__name__.startswith("Uncalibrated"):
                kernel = tfp.mcmc.MetropolisHastings(kernel)
            self._init_fn, self._sample_fn = _extract_kernel_functions(kernel)
        return init_params

    @property
    def model(self):
        return self._model

    @property
    def sample_field(self):
        return 'z'

    @property
    def default_fields(self):
        return ('z',)

    def get_diagnostics_str(self, state):
        """
        Given the current `state`, returns the diagnostics string to
        be added to progress bar for diagnostics purpose.
        """
        return ''

    def init(self, rng_key, num_warmup, init_params=None, model_args=(), model_kwargs={}):
        # non-vectorized
        if rng_key.ndim == 1:
            rng_key, rng_key_init_model = random.split(rng_key)
        # vectorized
        else:
            rng_key, rng_key_init_model = jnp.swapaxes(vmap(random.split)(rng_key), 0, 1)
        init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params)
        if self._potential_fn and init_params is None:
            raise ValueError('Valid value of `init_params` must be provided with'
                             ' `target_log_prob_fn`.')

        if rng_key.ndim == 1:
            init_state = self._init_fn(init_params, rng_key)
        else:
            # XXX it is safe to run hmc_init_fn under vmap despite that hmc_init_fn changes some
            # nonlocal variables: momentum_generator, wa_update, trajectory_len, max_treedepth,
            # wa_steps because those variables do not depend on traced args: init_params, rng_key.
            init_state = vmap(self._init_fn)(init_params, rng_key)
            sample_fn = vmap(self._sample_fn, in_axes=(0, None, None))
            self._sample_fn = sample_fn
        return init_state

    def postprocess_fn(self, args, kwargs):
        if self._postprocess_fn is None:
            return identity
        return self._postprocess_fn(*args, **kwargs)

    def sample(self, state, model_args, model_kwargs):
        """
        Run the kernel from the given :data:`~numpyro.contrib.tfp.mcmc.TFPKernelState`
        and return the resulting :data:`~numpyro.contrib.tfp.mcmc.TFPKernelState`.

        :param TFPKernelState state: Represents the current state.
        :param model_args: Arguments provided to the model.
        :param model_kwargs: Keyword arguments provided to the model.
        :return: Next `state` after running the kernel.
        """
        return self._sample_fn(state, model_args, model_kwargs)

__all__ = ['TFPKernel']
for _name, _Kernel in tfp.mcmc.__dict__.items():
    if not isinstance(_Kernel, type):
        continue
    if not issubclass(_Kernel, tfp.mcmc.TransitionKernel):
        continue
    if 'target_log_prob_fn' not in inspect.getfullargspec(_Kernel).args:
        continue

    _PyroKernel = TFPKernel[_Kernel]
    _PyroKernel.__module__ = __name__
    locals()[_name] = _PyroKernel

    _PyroKernel.__doc__ = '''
    Wraps `{}.{} <https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/mcmc/{}>`_
    with :class:`~numpyro.contrib.tfp.mcmc.TFPKernel`. The first argument `target_log_prob_fn`
    in TFP kernel construction is replaced by either `model` or `potential_fn`.
    '''.format(_Kernel.__module__, _Kernel.__name__, _Kernel.__name__)

    __all__.append(_name)

# Create sphinx documentation.
__doc__ = '\n\n'.join([

    '''
    {0}
    ----------------------------------------------------------------
    .. autoclass:: numpyro.contrib.tfp.mcmc.{0}
    '''.format(_name)
    for _name in __all__[:1] + sorted(__all__[1:])
])