import os
import itertools
import sys
import math
import numpy
import xml.etree.ElementTree as etree
from copy import copy
from collections import defaultdict
from datetime import date
try:
    from simtk.openmm import *
    from simtk.openmm.app import *
    import simtk.unit as units
    from simtk.openmm.app import element as elem
except:
    from openmm import *
    from openmm.app import *
    import openmm.unit as units
    from openmm.app import element as elem

"""
This is a modified 'simtk.openmm.app.forcefield.py'
to generate the amber prmtop file for the non-standard Residue.
For example, in case that the amino acid is covalented to the co-factor or ligand,
we cannot use the standard protein database (data/protein.ff14SB.xml).
"""

class AtomType (object):

    def __init__ (self, name, atomClass, mass, element):
        self.name      = name
        self.atomClass = atomClass
        self.mass      = mass
        self.element   = element

class AtomData (object):

    def __init__ (self, str_name, str_type, str_element, str_chg):

        self.at_name = str_name
        self.at_type = str_type
        self.at_element = str_element
        self.at_chg  = float(str_chg)
        self.bondedTo = []
        self.externalBonds = 0

class BondData (object):

    def __init__ (self, iatom, jatom):

        self.atom1 = iatom
        self.atom2 = jatom
        self.length = 0.0

        if iatom > jatom:
            self.atom2 = iatom
            self.atom1 = jatom

class ResidueData (object):
    """
     AMBER PRMTOP did not support virtual sites with massless.
    """
    def __init__ (self, res_name):

        self.name  = res_name
        self.atoms = []  # AtomData
        self.atomIndices = {}
        self.bonds = []
        self.externalBonds = []

    def addBond (self, atom1, atom2):
        if atom1 < atom2:
            self.bonds.append ( (atom1, atom2) )
        else:
            self.bonds.append ( (atom2, atom1) )
        self.atoms[atom1].bondedTo.append (atom2)
        self.atoms[atom2].bondedTo.append (atom1)

    def addBondByName (self, atom1_name, atom2_name):
        atom1 = self.atomIndices[atom1_name]
        atom2 = self.atomIndices[atom2_name]
        self.addBond (atom1, atom2)

    def addExternalBond (self, atom1):
        self.externalBonds.append (atom1)
        self.atoms[atom1].externalBonds += 1

    def addExternalBondByName (self, atom1_name):
        atom1 = self.atomIndices[atom1_name]
        self.addExternalBond (atom1)

class HarmonicBondData(object):
    """ a HarmonicBondForce. """
    def __init__ (self, type1, type2, length, k):
        self.types1 = type1
        self.types2 = type2
        if type1 > type2:
            self.types1 = type2
            self.types2 = type1
        self.length = length
        self.k      = k

class HarmonicAngleData(object):
    """ a HarmonicAngleForce. """
    def __init__ (self, type1, type2, type3, angle, k):
        self.types1 = type1
        self.types2 = type2
        self.types3 = type3
        if type1 > type3:
            self.types1 = type3
            self.types3 = type1
        self.angle  = angle
        self.k      = k

class PeriodicTorsionData (object):

    def __init__ (self):
        self.periodicity = []
        self.phase       = []
        self.k           = []

class PeriodicTorsion(object):

    def __init__(self, types):

        self.types1 = types[0]
        self.types2 = types[1]
        self.types3 = types[2]
        self.types4 = types[3]
        self.index_torsion = -1
        self.ordering    = 'default'

class MyForceFields (object):

    def __init__ (self, file_names): #pdb_file_name):
        # read filename

        self._atomTypes = {}
        self._residues  = {}
        self._atomClasses = {'':set()}
        self._harmonicBonds = []
        self._harmonicAngles = []
        self._propers   = {}
        self._impropers = {}
        self._nonbonds  = {}
        self._nonbonds_coulomb14scale = 0.8333333333333334
        self._nonbonds_lj14scale      = 0.5
        self._bondsForAtomType = defaultdict(set)
        self._unique_torsion_list = []
        self._unique_lj_list = []

        trees = []

        for xml_file_name in file_names:
            print ('file_names', xml_file_name)
            tree = etree.parse (xml_file_name)
            trees.append(tree)

        # Load the atom types
        for tree in trees:
            element = tree.getroot().find('AtomTypes')
            if element is not None:
                for at_type in element.findall('Type'):
                    typeName    = at_type.attrib['name']
                    atomClass   = at_type.attrib['class']
                    atomMass    = float(at_type.attrib['mass'])
                    atomElement = None
                    if 'element' in at_type.attrib:
                        atomElement = at_type.attrib['element']
                    self._atomTypes[typeName] = AtomType (typeName, atomClass,
                                                          atomMass, atomElement)

                    if atomClass not in self._atomClasses:
                        self._atomClasses[atomClass] = set()

                    self._atomClasses[atomClass].add(typeName)
                    self._atomClasses[''].add(typeName)

        # Load the Residue templates
        for tree in trees:
            element = tree.getroot().find('Residues')
            if element is not None:
                for residue in element.findall('Residue'):
                    resName     = residue.attrib['name']
                    resData     = ResidueData (resName)

                    #sum_charge = 0.0

                    for ia, atom in enumerate (residue.findall('Atom')):
                        atomName = atom.attrib['name']
                        typeName = atom.attrib['type']
                        charge   = atom.attrib['charge']

                     #   sum_charge += float(charge)

                        resData.atomIndices[atomName] = ia

                        atomData = AtomData(atomName,
                                            typeName,
                                            self._atomTypes[typeName].element,
                                            charge)

                        resData.atoms.append (atomData)

                    #if abs(sum_charge) > 0.001:
                    #    print (resName, sum_charge)

                    for bond in residue.findall('Bond'):
                        resData.addBondByName(bond.attrib['atomName1'],
                                              bond.attrib['atomName2'])

                    for bond in residue.findall('ExternalBond'):
                        resData.addExternalBondByName(bond.attrib['atomName'])

                    # Register Template
                    self._residues[resData.name] = resData

        # Load the HarmonicBondForce
        # OpenMM = 1/2 k (x - x0)**2 : E(kJ/mol/nm**2)
        # Amber  = k' (x-x0)**2      : E(kcal/mol/A**2)
        # k' = 1/2 k
        ene_conv = (units.kilojoule_per_mole/(units.nanometer*units.nanometer)
                             ).conversion_factor_to( units.kilocalorie_per_mole/(units.angstrom*units.angstrom))
        len_conv = units.nanometer.conversion_factor_to(units.angstrom)

        for tree in trees:
            element = tree.getroot().find('HarmonicBondForce')
            if element is not None:
                for bond in element.findall('Bond'):
                    type1 = bond.attrib['type1']
                    type2 = bond.attrib['type2']
                    bond_length = float(bond.attrib['length'])*len_conv
                    bond_k      = 0.5*float(bond.attrib['k'])*ene_conv
                    harmBond    = HarmonicBondData (type1, type2, bond_length, bond_k)
                    self._harmonicBonds.append ( harmBond )

        ene_conv = units.kilojoule_per_mole.conversion_factor_to (units.kilocalorie_per_mole)
        for tree in trees:
            element = tree.getroot().find('HarmonicAngleForce')
            if element is not None:
                for angle in element.findall('Angle'):
                    type1 = angle.attrib['type1']
                    type2 = angle.attrib['type2']
                    type3 = angle.attrib['type3']
                    ang_length = float(angle.attrib['angle'])
                    ang_k      = 0.5*float(angle.attrib['k'])*ene_conv
                    harmAngle  = HarmonicAngleData (type1, type2, type3,
                                                    ang_length, ang_k)
                    self._harmonicAngles.append ( harmAngle )

        ene_conv = units.kilojoule_per_mole.conversion_factor_to (units.kilocalorie_per_mole)
        for tree in trees:
            element = tree.getroot().find('PeriodicTorsionForce')
            if element is not None:

                ordering = 'default'
                if 'ordering' in element.attrib:
                    ordering = element.attrib['ordering']

                for proper in element.findall('Proper'):
                    types = []
                    for i in range(4):
                        suffix = str(i+1)
                        typeAttrib = 'type'+suffix
                        typeName   = proper.attrib[typeAttrib]
                        if typeName == '':
                            types.append('X') #self._atomClasses[''])
                        elif typeName not in self._atomTypes:
                            types.append(None)
                            print ('Proper ', typeName)
                        else:
                            types.append(typeName)

                    torsion = PeriodicTorsion(types)
                    index  = 1
                    torsionData = PeriodicTorsionData()

                    while 'phase%d'%index in proper.attrib:
                        periodicity = int(proper.attrib['periodicity%d'%index])
                        phase = float(proper.attrib['phase%d'%index])
                        k     = float(proper.attrib['k%d'%index])*ene_conv
                        torsionData.periodicity.append(periodicity)
                        torsionData.phase.append(phase)
                        torsionData.k.append(k)
                        index += 1

                    ladd = 1
                    for uniqueTorsion in self._unique_torsion_list:
                        if torsionData.periodicity == uniqueTorsion.periodicity and \
                           torsionData.phase       == uniqueTorsion.phase and \
                           torsionData.k           == uniqueTorsion.k:
                            ladd = 0
                            break

                    if ladd == 1:
                        self._unique_torsion_list.append (torsionData)

                    for ii in range (len(self._unique_torsion_list)):
                        uniqueTorsion = self._unique_torsion_list[ii]
                        if torsionData.periodicity == uniqueTorsion.periodicity and \
                           torsionData.phase       == uniqueTorsion.phase and \
                           torsionData.k           == uniqueTorsion.k:
                            torsion.index_torsion = ii

                    typeID  = 'PR_'
                    if types[0] == types[3]:
                        if types[1] < types[2]:
                            typeID += types[0]+'_'
                            typeID += types[1]+'_'
                            typeID += types[2]+'_'
                            typeID += types[3]
                        else:
                            typeID += types[3]+'_'
                            typeID += types[2]+'_'
                            typeID += types[1]+'_'
                            typeID += types[0]

                    elif types[0] < types[3]:
                        typeID += types[0]+'_'
                        typeID += types[1]+'_'
                        typeID += types[2]+'_'
                        typeID += types[3]
                    else:
                        typeID += types[3]+'_'
                        typeID += types[2]+'_'
                        typeID += types[1]+'_'
                        typeID += types[0]

                    self._propers[typeID] = torsion

                for improper in element.findall('Improper'):
                    types = []
                    improper_openmm2amber = [2,3,1,4]
                    for i in improper_openmm2amber:
                        suffix = str(i)
                        typeAttrib = 'type'+suffix
                        typeName   = improper.attrib[typeAttrib]
                        if typeName == '':
                            types.append('X')#self._atomClasses[''])
                        elif typeName not in self._atomTypes:
                            types.append(None)
                            print ('ImProper ', typeName)
                        else:
                            types.append(typeName)

                    torsion = PeriodicTorsion(types)
                    torsionData = PeriodicTorsionData()

                    torsion.ordering = ordering
                    periodicity = int(improper.attrib['periodicity1'])
                    phase       = float(improper.attrib['phase1'])
                    k           = float(improper.attrib['k1'])*ene_conv
                    torsionData.periodicity.append(periodicity)
                    torsionData.phase.append(phase)
                    torsionData.k.append(k)

                    ladd = 1
                    for uniqueTorsion in self._unique_torsion_list:
                        if torsionData.periodicity == uniqueTorsion.periodicity and \
                           torsionData.phase       == uniqueTorsion.phase and \
                           torsionData.k           == uniqueTorsion.k:
                            ladd = 0
                            break

                    if ladd == 1:
                        self._unique_torsion_list.append (torsionData)

                    for ii in range (len(self._unique_torsion_list)):
                        uniqueTorsion = self._unique_torsion_list[ii]
                        if torsionData.periodicity == uniqueTorsion.periodicity and \
                           torsionData.phase       == uniqueTorsion.phase and \
                           torsionData.k           == uniqueTorsion.k:
                            torsion.index_torsion = ii

                    if types[0] > types[1]:
                        (types[0], types[1]) = (types[1], types[0])

                    typeID  = 'IM_'
                    typeID += types[0]+'_'
                    typeID += types[1]+'_'
                    typeID += types[2]+'_'
                    typeID += types[3]

                    self._impropers[typeID] = torsion

        # Eunit: kJ/mol --> kcal/mol
        # Lunit: nm --> A
        ene_conv = units.kilojoules_per_mole.conversion_factor_to(units.kilocalories_per_mole)
        len_conv = units.nanometers.conversion_factor_to(units.angstrom)
        # Load NonbondedForce
        for tree in trees:
            element = tree.getroot().find('NonbondedForce')
            if element is not None:

                if 'coulomb14scale' in element.attrib:
                    self._nonbonds_coulomb14scale = float(element.attrib['coulomb14scale'])
                if 'lj14scale' in element.attrib:
                    self._nonbonds_lj14scale      = float(element.attrib['lj14scale'])
                for atom in element.findall('Atom'):
                    typeName = atom.attrib['type']

                    epsilon  = float(atom.attrib['epsilon'])*ene_conv
                    sigma    = float(atom.attrib['sigma'])*len_conv

                    self._nonbonds[typeName] = [epsilon, sigma]


__all__ = ["translation_rule"]

import numpy as np
from jax.lib import xla_client

from . import jax_finufft

for _name, _value in jax_finufft.registrations().items():
    xla_client.register_cpu_custom_call_target(_name, _value)

xops = xla_client.ops

def translation_rule(
    ctx, avals_in, avals_out, source, *points, output_shape, iflag, eps
):
    ndim = len(points)
    assert 1 <= ndim <= 3

    c = ctx.builder
    source_shape_info = c.get_shape(source)
    points_shape_info = list(map(c.get_shape, points))

    # Check supported and consistent dtypes
    source_dtype = source_shape_info.element_type()
    single = source_dtype == np.csingle and all(
        x.element_type() == np.single for x in points_shape_info
    )
    double = source_dtype == np.cdouble and all(
        x.element_type() == np.double for x in points_shape_info
    )
    assert single or double
    suffix = "f" if source_dtype == np.csingle else ""

    # Check shapes
    source_shape = source_shape_info.dimensions()
    points_shape = tuple(x.dimensions() for x in points_shape_info)
    n_tot = source_shape[0]
    n_transf = source_shape[1]
    n_j = points_shape[0][1]
    if output_shape is None:
        op_name = f"nufft{ndim}d2{suffix}".encode("ascii")
        n_k = np.array(source_shape[2:], dtype=np.int64)
        full_output_shape = source_shape[:2] + (n_j,)
    else:
        type_ = 1
        op_name = f"nufft{ndim}d1{suffix}".encode("ascii")
        n_k = np.array(output_shape, dtype=np.int64)
        full_output_shape = source_shape[:2] + tuple(output_shape)

    # The backend expects the output shape in Fortran order so we'll just
    # fake it here, by sending in n_k and x in the reverse order.
    n_k_full = np.zeros(3, dtype=np.int64)
    n_k_full[:ndim] = n_k[::-1]

    # Dispatch to the right op
    desc = getattr(jax_finufft, f"build_descriptor{suffix}")(
        eps, iflag, n_tot, n_transf, n_j, *n_k_full
    )

    return [
        xops.CustomCallWithLayout(
            c,
            op_name,
            # The inputs:
            operands=(
                xops.ConstantLiteral(c, np.frombuffer(desc, dtype=np.uint8)),
                source,
                *points[::-1],  # Reverse order because backend uses Fortran order
            ),
            # The input shapes:
            operand_shapes_with_layout=(
                xla_client.Shape.array_shape(np.dtype(np.uint8), (len(desc),), (0,)),
                xla_client.Shape.array_shape(
                    source_dtype,
                    source_shape,
                    tuple(range(len(source_shape) - 1, -1, -1)),
                ),
            )
            + tuple(
                xla_client.Shape.array_shape(
                    x.element_type(),
                    x.dimensions(),
                    tuple(range(len(x.dimensions()) - 1, -1, -1)),
                )
                for x in points_shape_info[::-1]  # Reverse order, again
            ),
            # The output shapes:
            shape_with_layout=xla_client.Shape.array_shape(
                source_dtype,
                full_output_shape,
                tuple(range(len(full_output_shape) - 1, -1, -1)),
            ),
        )
    ]


# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

# Simple background function.  Assumes background has form approx 1/sqrt(m2q).
# Note. bg_param assumed a 1 dalton binning.  Mult by binsize for other sizes.
#

# standard imports
import numpy as np
import matplotlib.pyplot as plt

# custom imports
import apt_fileio
import plotting_stuff
import initElements_P3

import peak_param_determination as ppd

from histogram_functions import bin_dat

plt.close('all')

# Read in data
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\180821_GaN_A71\R20_07094-v03.epos"
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\R20_07148-v01.epos" # Mg doped
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\R20_07148-v01_vbmq_corr.epos"
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\R20_07247.epos"
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\R20_07248-v01.epos"
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\R20_07249-v01.epos"
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\R20_07250-v01.epos"
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\190421_AlGaN50p7_A83\R20_07209-v01.epos"
#
fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\181210_D315_A74\R20_07167-v03.epos"
#fn = r"\\cfs2w.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\181210_D315_A74\R20_07148-v02.epos"
#fn = r"\\cfs2w.campus.nist.gov\647\NIST_Projects\EUV_APT_IMS\BWC\GaN epos files\181204_InGaNQW_A73\R20_07144-v02.epos"

fn = fn[:-5]+'_vbm_corr.epos'
epos = apt_fileio.read_epos_numpy(fn)
#epos = epos[epos.size//2:-1]

# Plot m2q vs event index and show the current ROI selection
roi_event_idxs = np.arange(5000,epos.size-10000)

#roi_event_idxs = np.arange(epos.size)
ax = plotting_stuff.plot_m2q_vs_time(epos['m2q'],epos,fig_idx=1)
ax.plot(roi_event_idxs[0]*np.ones(2),[0,1200],'--k')
ax.plot(roi_event_idxs[-1]*np.ones(2),[0,1200],'--k')
ax.set_title('roi selected to start analysis')
epos = epos[roi_event_idxs]

# Compute some extra information from epos information
LASER_REP_RATE = 10000.0
wall_time = np.cumsum(epos['pslep'])/LASER_REP_RATE
pulse_idx = np.arange(0,epos.size)
isSingle = np.nonzero(epos['ipp'] == 1)

# Define peaks to range
ed = initElements_P3.initElements()

#                            N      Ga       Da
# Define possible peaks
pk_data =   np.array(    [  (1,     0,        ed['N'].isotopes[14][0]/2),
                            (1,     0,        ed['N'].isotopes[14][0]/1),
                            (1,     0,        ed['N'].isotopes[15][0]/1),
                            (1,     0,        ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                            (0,     1,        ed['Ga'].isotopes[69][0]/3),
                            (0,     1,        ed['Ga'].isotopes[71][0]/3),
                            (2,     0,        ed['N'].isotopes[14][0]*2),
                            (2,     0,        ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
                            (2,     0,        ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
                            (0,     1,        ed['Ga'].isotopes[69][0]/2),
                            (0,     1,        ed['Ga'].isotopes[71][0]/2),
                            (1,     1,        (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
                            (3,     0,        ed['N'].isotopes[14][0]*3),
                            (1,     1,        (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
                            (3,     1,        (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
                            (3,     1,        (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
                            (0,     1,        ed['Ga'].isotopes[69][0]),
                            (0,     1,        ed['Ga'].isotopes[71][0]),
                            (0,     1,        ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
                            ],
                            dtype=[('N','i4'),('Ga','i4'),('m2q','f4')] )

#                            N      Ga      In  Da
## Define possible peaks
#
#pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
#                            (1,     0,      0,  ed['N'].isotopes[14][0]/1),
#                            (1,     0,      0,  ed['N'].isotopes[15][0]/1),
#                            (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
#                            (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
#                            (2,     0,      0,  ed['N'].isotopes[14][0]*2),
#                            (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
#                            (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
#                            (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
#                            (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
#                            (3,     0,      0,  ed['N'].isotopes[14][0]*3),
#                            (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
#                            (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
#                            (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
#                            (0,     1,      0,  ed['Ga'].isotopes[69][0]),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0]),
#                            (0,     0,      1,  ed['In'].isotopes[113][0]/2),
#                            (0,     0,      1,  ed['In'].isotopes[115][0]/2)
#                            ],
#                            dtype=[('N','i4'),('Ga','i4'),('In','i4'),('m2q','f4')] )

##                            N      Ga      Al  Da
## Define possible peaks
#
#pk_data =   np.array(    [  (1,     0,      0,  ed['N'].isotopes[14][0]/2),
#                            (0,     0,      1,  ed['Al'].isotopes[27][0]/3),
#                            (0,     0,      1,  ed['Al'].isotopes[27][0]/2),
#                            (1,     0,      0,  ed['N'].isotopes[14][0]/1),
#                            (1,     0,      0,  ed['N'].isotopes[15][0]/1),
#                            (1,     0,      0,  ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
#                            (1,     0,      1,  (ed['N'].isotopes[14][0]+ed['Al'].isotopes[27][0])/2),
#                            (0,     1,      0,  ed['Ga'].isotopes[69][0]/3),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]/3),
#                            (0,     0,      1,  ed['Al'].isotopes[27][0]/1),
#                            (2,     0,      0,  ed['N'].isotopes[14][0]*2),
#                            (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]),
#                            (2,     0,      0,  ed['N'].isotopes[14][0]+ed['N'].isotopes[15][0]+ed['H'].isotopes[1][0]),
#                            (0,     1,      0,  ed['Ga'].isotopes[69][0]/2),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]/2),
#                            (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[69][0])/2),
#                            (3,     0,      0,  ed['N'].isotopes[14][0]*3),
#                            (1,     1,      0,  (ed['N'].isotopes[14][0] + ed['Ga'].isotopes[71][0])/2),
#                            (3,     1,      0,  (ed['Ga'].isotopes[69][0]+3*ed['N'].isotopes[14][0])/2),
#                            (3,     1,      0,  (ed['Ga'].isotopes[71][0]+3*ed['N'].isotopes[14][0])/2),
#                            (0,     1,      0,  ed['Ga'].isotopes[69][0]),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]),
#                            (0,     1,      0,  ed['Ga'].isotopes[71][0]+ed['H'].isotopes[1][0])
#                            ],
#                            dtype=[('N','i4'),('Ga','i4'),('Al','i4'),('m2q','f4')] )

# Define which peaks to use for CSR calcs
Ga1p_m2qs = [ed['Ga'].isotopes[69][0], ed['Ga'].isotopes[71][0]]
Ga2p_m2qs = [ed['Ga'].isotopes[69][0]/2, ed['Ga'].isotopes[71][0]/2]

Ga1p_idxs = [np.argmin(np.abs(m2q-pk_data['m2q'])) for m2q in Ga1p_m2qs]
Ga2p_idxs = [np.argmin(np.abs(m2q-pk_data['m2q'])) for m2q in Ga2p_m2qs]

# Range the peaks
pk_params = ppd.get_peak_ranges(epos,pk_data['m2q'],peak_height_fraction=0.1)

# Determine the global background
#glob_bg_param = ppd.fit_uncorr_bg(epos['m2q'],fit_roi=[3.5,6.5])
bg_rois=[[3.5,6.5],[90,110]]
glob_bg_param = ppd.get_glob_bg(epos['m2q'],rois=bg_rois)

# Count the peaks, local bg, and global bg
cts = ppd.do_counting(epos,pk_params,glob_bg_param)

# Test for peak S/N and throw out craptastic peaks
B = np.max(np.c_[cts['local_bg'][:,None],cts['global_bg'][:,None]],1)[:,None]
T = cts['total'][:,None]
S = T-B
std_S = np.sqrt(T+B)
# Make up a threshold for peak detection... for the most part this won't matter
# since weak peaks don't contribute to stoichiometry much... except for Mg!
is_peak = S>2*np.sqrt(2*B)
for idx, ct in enumerate(cts):
    if not is_peak[idx]:
        for i in np.arange(len(ct)):
            ct[i] = 0

# Calculate compositions
compositions = ppd.do_composition(pk_data,cts)
ppd.pretty_print_compositions(compositions,pk_data)

print('Total Ranged Ions: '+str(np.sum(cts['total'])))
print('Total Ranged Local Background Ions: '+str(np.sum(cts['local_bg'])))
print('Total Ranged Global Background Ions: '+str(np.sum(cts['global_bg'])))
print('Total Ions: '+str(epos.size))

print('Overall CSR (no bg)    : '+str(np.sum(cts['total'][Ga2p_idxs])/np.sum(cts['total'][Ga1p_idxs])))
print('Overall CSR (local bg) : '+str((np.sum(cts['total'][Ga2p_idxs])-np.sum(cts['local_bg'][Ga2p_idxs]))/(np.sum(cts['total'][Ga1p_idxs])-np.sum(cts['local_bg'][Ga1p_idxs]))))
print('Overall CSR (global bg): '+str((np.sum(cts['total'][Ga2p_idxs])-np.sum(cts['global_bg'][Ga2p_idxs]))/(np.sum(cts['total'][Ga1p_idxs])-np.sum(cts['global_bg'][Ga1p_idxs]))))

# Plot all the things
xs, ys = bin_dat(epos['m2q'],user_roi=[0.5, 120],isBinAligned=True)
#ys_sm = ppd.do_smooth_with_gaussian(ys,30)
ys_sm = ppd.moving_average(ys,30)

glob_bg = ppd.physics_bg(xs,glob_bg_param)

fig = plt.figure(num=100)
fig.clear()
ax = fig.gca()

ax.plot(xs,ys_sm,label='hist')
ax.plot(xs,glob_bg,label='global bg')

ax.set(xlabel='m/z (Da)', ylabel='counts')
ax.grid()
fig.tight_layout()
fig.canvas.manager.window.raise_()
ax.set_yscale('log')
ax.legend()

for idx,pk_param in enumerate(pk_params):

    if is_peak[idx]:
        ax.plot(np.array([1,1])*pk_param['pre_rng'] ,np.array([0.5,(pk_param['amp']+pk_param['off'])]),'k--')
        ax.plot(np.array([1,1])*pk_param['post_rng'] ,np.array([0.5,(pk_param['amp']+pk_param['off'])]),'k--')
        ax.plot(np.array([1,1])*pk_param['pre_bg_rng'] ,np.array([0.5,(pk_param['amp']+pk_param['off'])]),'m--')
        ax.plot(np.array([1,1])*pk_param['post_bg_rng'] ,np.array([0.5,(pk_param['amp']+pk_param['off'])]),'m--')

        ax.plot(np.array([pk_param['pre_bg_rng'],pk_param['post_bg_rng']]) ,np.ones(2)*pk_param['loc_bg'],'g--')
    else:
        ax.plot(np.array([1,1])*pk_param['x0_mean_shift'] ,np.array([0.5,(pk_param['amp']+pk_param['off'])]),'r--')

for roi in bg_rois:
    xbox = np.array([roi[0],roi[0],roi[1],roi[1]])
    ybox = np.array([0.1,np.max(ys_sm)/10,np.max(ys_sm)/10,0.1])

    ax.fill(xbox,ybox, 'b', alpha=0.2)

plt.pause(0.1)

#### END BASIC ANALYSIS ####

#### START EXPLORATORY ANALYSIS ####

# Slice and dice the data in wall_time
idxs_list = []
CHUNK_SIZE = 32000
s_idx = 0
while s_idx<epos.size:
    e_idx = np.min([epos.size,s_idx+CHUNK_SIZE])
    idxs_list.append([s_idx,e_idx])
    s_idx = e_idx

# Count and compositions
csr = np.full(len(idxs_list),-1.0)
Ga_comp = np.full(len(idxs_list),-1.0)
Ga_comp_std = np.full(len(idxs_list),-1.0)
Ga_comp_glob = np.full(len(idxs_list),-1.0)
Ga_comp_std_glob = np.full(len(idxs_list),-1.0)

keys = list(pk_data.dtype.fields.keys())
keys.remove('m2q')
Ga_idx = keys.index('Ga')
for loop_idx, idxs in enumerate(idxs_list):

    sub_epos = epos[idxs[0]:idxs[1]]

    plotting_stuff.plot_histo(sub_epos['m2q'],321,user_xlim=[0,275],user_bin_width=0.1,user_label=loop_idx)
    plt.waitforbuttonpress()
#    r = np.sqrt(sub_epos['x_det']**2+sub_epos['y_det']**2)
#    sub_idxs = np.nonzero(r>=0)
#    cts = ppd.do_counting(sub_epos[sub_idxs],pk_params,glob_bg_param)
#
#    tot_bg_ct = epos['m2q'][(epos['m2q']>=80) & (epos['m2q']<=120)].size
#    sub_bg_ct = sub_epos['m2q'][(sub_epos['m2q']>=80) & (sub_epos['m2q']<=120)].size
#
#    loc_bg_param = glob_bg_param*sub_bg_ct/tot_bg_ct
    glob_bg_param_chunk = ppd.get_glob_bg(sub_epos['m2q'],rois=bg_rois)

    cts = ppd.do_counting(sub_epos,pk_params,glob_bg_param_chunk)
    Ga_comp[loop_idx] = ppd.do_composition(pk_data,cts)[0][0][Ga_idx]
    Ga_comp_std[loop_idx] = ppd.do_composition(pk_data,cts)[0][1][Ga_idx]

    Ga_comp_glob[loop_idx] = ppd.do_composition(pk_data,cts)[2][0][Ga_idx]
    Ga_comp_std_glob[loop_idx] = ppd.do_composition(pk_data,cts)[2][1][Ga_idx]

    csr[loop_idx] = np.sum(cts['total'][Ga2p_idxs])/np.sum(cts['total'][Ga1p_idxs])
    print('Total Ranged Ions: '+str(np.sum(cts['total'])))
    print('Total Ranged Global Background Ions: '+str(np.sum(cts['global_bg'])))
    print('Total Ions: '+str(sub_epos.size))

    compositions = ppd.do_composition(pk_data,cts)
    ppd.pretty_print_compositions(compositions,pk_data)
#
#
#    print(sub_epos['m2q'][(sub_epos['m2q']>=80) & (sub_epos['m2q']<=120)].size)
#    print(np.sum(cts['total']))
#    print(sub_epos['m2q'][(sub_epos['m2q']>=80) & (sub_epos['m2q']<=120)].size/np.sum(cts['total']))
#

fig = plt.figure(num=101)
fig.clear()
ax = fig.gca()
ax.errorbar((np.arange(Ga_comp.size)+0.5)*CHUNK_SIZE,Ga_comp,yerr=Ga_comp_std,fmt='.',capsize=4,label='chunk based (wall time)')
ax.errorbar((np.arange(Ga_comp.size)+0.5)*CHUNK_SIZE,Ga_comp_glob,yerr=Ga_comp_std_glob,fmt='.',capsize=4,label='chunk based (wall time)_glob')
ax.set(xlabel='ion idx', ylabel='Ga %')
ax.legend()

fig = plt.figure(num=102)
fig.clear()
ax = fig.gca()
ax.errorbar(csr,Ga_comp,yerr=Ga_comp_std,fmt='.',capsize=4,label='by wall time')

ax.plot([np.min(csr),np.max(csr)],[0.5,0.5],'k--')

ax.set(xlabel='CSR', ylabel='Ga %', ylim=[0, 1], xlim=[1e-2,2])

ax.set_title('by wall time')
ax.set_xscale('log')

ax.legend()
ax.grid()
fig.tight_layout()
fig.canvas.manager.window.raise_()

import colorcet as cc

def extents(f):
    delta = f[1] - f[0]
    return [f[0] - delta/2, f[-1] + delta/2]

def create_histogram(xs,ys,x_roi=None,y_roi=None):
    num_x = 128
    num_y = num_x
    N,x_edges,y_edges = np.histogram2d(xs,ys,bins=[num_x,num_y],range=[x_roi,y_roi],density=False)
    return (N,x_edges,y_edges)

fig = plt.figure(num=301)
plt.clf()
ax = fig.gca()

sel_idxs = np.arange(epos.size)
sel_idxs = np.where((epos['m2q']>(0.9+0)) & (epos['m2q']<(1.2+0)))

pk_data

N,x_edges,y_edges = create_histogram(epos['x_det'][sel_idxs],epos['y_det'][sel_idxs],
                                     x_roi=[-35,35],y_roi=[-35,35])
#ax.imshow(np.log10(1+1*np.transpose(N)), aspect='auto',
ax.imshow(np.transpose(N), aspect='auto',
           extent=extents(x_edges) + extents(y_edges), origin='lower', cmap=cc.cm.CET_L8,
           interpolation='nearest')
ax.set_aspect('equal', 'box')

ax.set(xlabel='det_x')
ax.set(ylabel='det_y')

keys = list(pk_data.dtype.fields.keys())
keys.remove('m2q')

for k in keys:
    k_pks = np.where(pk_data[k]>0)[0]
    sel_idxs = np.zeros(0,dtype='int64')
    for pk in k_pks:
        ev_idxs = np.where((epos['m2q']>pk_params['pre_rng'][pk]) & (epos['m2q']<pk_params['post_rng'][pk]))[0]
        sel_idxs = np.concatenate((sel_idxs,ev_idxs))

    fig = plt.figure()
    plt.clf()
    ax = fig.gca()

    N,x_edges,y_edges = create_histogram(epos['x_det'][sel_idxs],epos['y_det'][sel_idxs],
                                         x_roi=[-35,35],y_roi=[-35,35])
    #ax.imshow(np.log10(1+1*np.transpose(N)), aspect='auto',
    ax.imshow(np.transpose(N), aspect='auto',
               extent=extents(x_edges) + extents(y_edges), origin='lower', cmap=cc.cm.CET_L8,
               interpolation='nearest')
    ax.set_aspect('equal', 'box')

    ax.set(xlabel='det_x')
    ax.set(ylabel='det_y')
    ax.set_title(k)

# All ranged ions

keys = list(pk_data.dtype.fields.keys())
keys.remove('m2q')

sel_idxs = np.zeros(0,dtype='int64')

for k in keys:
    k_pks = np.where(pk_data[k]>0)[0]
    for pk in k_pks:
        ev_idxs = np.where((epos['m2q']>pk_params['pre_rng'][pk]) & (epos['m2q']<pk_params['post_rng'][pk]))[0]
        sel_idxs = np.concatenate((sel_idxs,ev_idxs))

fig = plt.figure(321)
plt.clf()
ax = fig.gca()

N,x_edges,y_edges = create_histogram(0+epos['x_det'][sel_idxs],0+epos['y_det'][sel_idxs],
                                     x_roi=[-35,35],y_roi=[-35,35])
#ax.imshow(np.log10(1+1*np.transpose(N)), aspect='auto',
met = N

met[met==0] = 100

met = 1/(met+3)

ax.imshow(np.transpose(met), aspect='auto',
           extent=extents(x_edges) + extents(y_edges), origin='lower', cmap=cc.cm.CET_R2,
           interpolation='nearest')
ax.set_aspect('equal', 'box')

ax.set(xlabel='det_x')
ax.set(ylabel='det_y')

cx = 0.5*(x_edges[1:]+x_edges[:-1])
cy = 0.5*(y_edges[1:]+y_edges[:-1])

CX,CY = np.meshgrid(cx,cy)

def mean_shift(xs,ys):
    radius = 5

    x_curr = 0
    y_curr = 0

    N_LOOPS = 64

    xi = np.zeros(N_LOOPS)
    yi = np.zeros(N_LOOPS)

    for i in np.arange(N_LOOPS):
        x_prev = x_curr
        y_prev = y_curr

        idxs = np.where(((xs-x_curr)**2+(ys-y_curr)**2) <= radius**2)
#        print(idxs)

        x_q = np.mean(xs[idxs])
        y_q = np.mean(ys[idxs])

        dx = x_q-x_prev
        dy = y_q-y_prev

        x_curr = x_prev-dx
        y_curr = y_prev-dy

        if np.sqrt((x_curr-x_prev)**2 + (y_curr-y_prev)**2) < (radius*1e-2):
#            print('iter  B #',i,'    ',x_curr,y_curr)
#            print(i)
            break
#        else:
#            print('iter NB #',i,'    ',x_curr,y_curr)
#        print('iter #',i,'    ',x_curr,y_curr)
        xi[i] = x_curr
        yi[i] = y_curr

    return xi[:i], yi[:i]

xc,yc = mean_shift(epos['x_det'][sel_idxs],epos['y_det'][sel_idxs])

ax.plot(xc,yc,'-o')

# Slice and dice the data in detector space (polar)
idxs_list = []
#STEP = 2

r = np.sqrt(np.square(epos['x_det']-xc[-1])+np.square(epos['y_det']-yc[-1]))

R_DET_MAX = 28
R_C = np.sqrt(xc[-1]**2+yc[-1]**2)

R_MAX = R_DET_MAX-R_C

r_edges = np.sqrt(np.linspace(0, R_MAX**2, 6))
r_centers = (r_edges[:-1]+r_edges[1:])/2

for i in np.arange(r_edges.size-1):
    idxs = np.where((r>r_edges[i]) & (r<=r_edges[i+1]))[0]
    idxs_list.append(idxs)

#for rq in np.arange(2,R_MAX,STEP):
#    idxs = np.where((r>rq) & (r<=(rq+STEP)))[0]
#    idxs_list.append(idxs)

# Count and compositions
csr = np.full(len(idxs_list),-1.0)
Ga_comp = np.full(len(idxs_list),-1.0)
Ga_comp_std = np.full(len(idxs_list),-1.0)
Ga_comp_glob = np.full(len(idxs_list),-1.0)
Ga_comp_std_glob = np.full(len(idxs_list),-1.0)

Ga_idx = keys.index('Ga')
for loop_idx, idxs in enumerate(idxs_list):
    sub_epos = epos[idxs]

#    tot_bg_ct = epos['m2q'][(epos['m2q']>=80) & (epos['m2q']<=120)].size
#    sub_bg_ct = sub_epos['m2q'][(sub_epos['m2q']>=80) & (sub_epos['m2q']<=120)].size

#    loc_bg_param = glob_bg_param*sub_bg_ct/tot_bg_ct

    glob_bg_param_chunk = ppd.get_glob_bg(sub_epos['m2q'],rois=bg_rois)

    cts = ppd.do_counting(sub_epos,pk_params,glob_bg_param_chunk)

    csr[loop_idx] = np.sum(cts['total'][Ga2p_idxs])/np.sum(cts['total'][Ga1p_idxs])
    Ga_comp[loop_idx] = ppd.do_composition(pk_data,cts)[0][0][Ga_idx]
    Ga_comp_std[loop_idx] = ppd.do_composition(pk_data,cts)[0][1][Ga_idx]

    Ga_comp_glob[loop_idx] = ppd.do_composition(pk_data,cts)[2][0][Ga_idx]
    Ga_comp_std_glob[loop_idx] = ppd.do_composition(pk_data,cts)[2][1][Ga_idx]

    compositions = ppd.do_composition(pk_data,cts)
    ppd.pretty_print_compositions(compositions,pk_data)
    print('COUNTS IN CHUNK: ',np.sum(cts['total']))

fig = plt.figure(num=201)
fig.clear()
ax = fig.gca()
#ax.errorbar((np.arange(Ga_comp.size)+0.5)*STEP,Ga_comp,yerr=Ga_comp_std,fmt='.',capsize=4,label='Ga %')
#ax.errorbar((np.arange(Ga_comp.size)+0.5)*STEP,Ga_comp_glob,yerr=Ga_comp_std_glob,fmt='.',capsize=4,label='glob')
ax.errorbar(r_centers,Ga_comp,yerr=Ga_comp_std,fmt='.',capsize=4,label='Ga %')
ax.errorbar(r_centers,Ga_comp_glob,yerr=Ga_comp_std_glob,fmt='.',capsize=4,label='glob')

ax.set(xlabel='radius', ylabel='Ga %')

ax_twin= ax.twinx()
#ax_twin.plot((np.arange(Ga_comp.size)+0.5)*STEP,csr,'s',color='r',label='Ga CSR')
ax_twin.plot(r_centers,csr,'s',color='r',label='Ga CSR')

ax.legend()
ax_twin.legend(loc=7)
ax_twin.set(ylabel='Ga CSR')

fig = plt.figure(num=202)
fig.clear()
ax = fig.gca()

ax.errorbar(csr,Ga_comp,yerr=Ga_comp_std,fmt='.',capsize=4,label='det based (radial)')
ax.plot([np.min(csr),np.max(csr)],[0.5,0.5],'k--')

ax.set(xlabel='CSR', ylabel='Ga %', ylim=[0, 1], xlim=[1e-2,2])
ax.legend()
ax.set_title('by radius')
ax.set_xscale('log')
ax.grid()
fig.tight_layout()
fig.canvas.manager.window.raise_()

for i in np.arange(csr.size):
    print(csr[i],'\t',Ga_comp[i],'\t',Ga_comp_std[i])

# chop data up by Radius AND time
# Aim for 50 k ions per time chunk
# Aim for 5 radial chunks

es2cs = lambda es : (es[:-1]+es[1:])/2.0

#idxs_list = []

N_time_chunks = int(np.floor(epos.size/256000))
#N_time_chunks = 4

N_events_per_time_chunk = epos.size//N_time_chunks
time_chunk_edges = np.arange(N_time_chunks+1)*N_events_per_time_chunk
time_chunk_centers = es2cs(time_chunk_edges)

R_DET_MAX = 28
R_C = np.sqrt(xc[-1]**2+yc[-1]**2)
R_MAX = R_DET_MAX-R_C

N_ann_chunks = 3
r_edges = np.sqrt(np.linspace(0, R_MAX**2, N_ann_chunks+1))
r_centers = es2cs(r_edges)

csr = np.full([N_time_chunks,N_ann_chunks],-1.0)
Ga_comp = np.full([N_time_chunks,N_ann_chunks],-1.0)
Ga_comp_std = np.full([N_time_chunks,N_ann_chunks],-1.0)
Ga_comp_glob = np.full([N_time_chunks,N_ann_chunks],-1.0)
Ga_comp_std_glob = np.full([N_time_chunks,N_ann_chunks],-1.0)
tot_cts = np.full([N_time_chunks,N_ann_chunks],-1.0)

for t_idx in np.arange(N_time_chunks):
    sub_epos = epos[time_chunk_edges[t_idx]:time_chunk_edges[t_idx+1]]
    r = np.sqrt(np.square(sub_epos['x_det']-xc[-1])+np.square(sub_epos['y_det']-yc[-1]))
    for a_idx in np.arange(N_ann_chunks):

        idxs = np.where((r>r_edges[a_idx]) & (r<=r_edges[a_idx+1]))[0]

        subsubepos = sub_epos[idxs]

        glob_bg_param_chunk = ppd.get_glob_bg(subsubepos['m2q'],rois=bg_rois)

        cts = ppd.do_counting(subsubepos,pk_params,glob_bg_param_chunk)

        csr[t_idx,a_idx] = np.sum(cts['total'][Ga2p_idxs])/np.sum(cts['total'][Ga1p_idxs])
        Ga_comp[t_idx,a_idx] = ppd.do_composition(pk_data,cts)[0][0][Ga_idx]
        Ga_comp_std[t_idx,a_idx] = ppd.do_composition(pk_data,cts)[0][1][Ga_idx]

        Ga_comp_glob[t_idx,a_idx] = ppd.do_composition(pk_data,cts)[2][0][Ga_idx]
        Ga_comp_std_glob[t_idx,a_idx] = ppd.do_composition(pk_data,cts)[2][1][Ga_idx]

        compositions = ppd.do_composition(pk_data,cts)
        ppd.pretty_print_compositions(compositions,pk_data)
        print('COUNTS IN CHUNK: ',np.sum(cts['total']))
        tot_cts[t_idx,a_idx] = np.sum(cts['total'])

fig = plt.figure(num=402)
fig.clear()
ax = fig.gca()

#ax.errorbar(csr.flatten(),Ga_comp.flatten(),yerr=Ga_comp_std.flatten(),fmt='.',capsize=4,label='det based (radial)')
ax.errorbar(csr.flatten(),Ga_comp_glob.flatten(),yerr=Ga_comp_std_glob.flatten(),fmt='.',capsize=4,label='det based (radial)')
ax.plot([np.min(csr),np.max(csr)],[0.5,0.5],'k--')

ax.set(xlabel='CSR', ylabel='Ga %', ylim=[0, 1], xlim=[5e-3,5])
ax.legend()
ax.set_title('by radius')
ax.set_xscale('log')
ax.grid()
fig.tight_layout()
fig.canvas.manager.window.raise_()

fig = plt.figure(num=405)
fig.clear()
ax = fig.gca()
ax.imshow(csr.T,
          extent=extents(time_chunk_edges) + [r_edges[0], r_edges[-1]+np.diff(r_edges[-2:])],
          aspect='auto',
           origin='lower')
ax.set(xlabel='ev idx')
ax.set(ylabel='rad')

for i in np.arange(csr.size):
    print(csr.flatten()[i],'\t',Ga_comp_glob.flatten()[i],'\t',Ga_comp_std_glob.flatten()[i])


__all__ = ['FixedTarget']

import numpy as np

from pylightcurve.errors import *
from pylightcurve.__databases__ import plc_data
from pylightcurve.analysis.curve_fit import curve_fit
from pylightcurve.spacetime.angles import arccos, _request_angle

class _Target:

    def __init__(self, ra, dec):

        _request_angle(ra)
        _request_angle(dec)

        if 270 > dec.deg() > 90:
            raise PyLCInputError('Declination must within -90, 90 degrees')

        self.dec = dec
        self.ra = ra

        self.coord = '{0} {1}'.format(self.ra.hms(), self.dec.dms_coord())

class FixedTarget(_Target):

    def __init__(self, ra, dec):

        _Target.__init__(self, ra, dec)

    def convert_to_bjd_tdb(self, time, time_format):

        if isinstance(time, float):

            if time_format in ['BJD_TDB', 'BJD_TT']:
                return time
            elif time_format == 'JD_UTC':
                return self._bjd_tdb(None, time)
            elif time_format == 'MJD_UTC':
                return self._bjd_tdb(None, time + 2400000.5)
            elif time_format == 'BJD_UTC':
                return self._bjd_tdb(None, curve_fit(self._bjd_utc, [0], [time], p0=[time])[0][0])
            elif time_format in ['HJD_TDB', 'HJD_TT']:
                return self._bjd_tdb(None, curve_fit(self._hjd_tdb, [0], [time], p0=[time])[0][0])
            elif time_format == 'HJD_UTC':
                return self._bjd_tdb(None, curve_fit(self._hjd_utc, [0], [time], p0=[time])[0][0])
            else:
                raise PyLCInputError(
                    'Not valid time format. Available formats: JD_UTC, MJD_UTC, BJD_UTC, BJD_TDB, BJD_TT, '
                    'HJD_UTC, HJD_BJD, HJD_TT')

        else:
            try:
                time = np.array(time, dtype=float)

                if time_format in ['BJD_TDB', 'BJD_TT']:
                    return time
                elif time_format == 'JD_UTC':
                    return np.array([self._bjd_tdb(None, ff) for ff in time])
                elif time_format == 'MJD_UTC':
                    return np.array([self._bjd_tdb(None, ff + 2400000.5) for ff in time])
                elif time_format == 'BJD_UTC':
                    return np.array([self._bjd_tdb(None, curve_fit(self._bjd_utc, [0], [ff], p0=[ff])[0][0])
                                     for ff in time])
                elif time_format in ['HJD_TDB', 'HJD_TT']:
                    return np.array([self._bjd_tdb(None, curve_fit(self._hjd_tdb, [0], [ff], p0=[ff])[0][0])
                                     for ff in time])
                elif time_format == 'HJD_UTC':
                    return np.array([self._bjd_tdb(None, curve_fit(self._hjd_utc, [0], [ff], p0=[ff])[0][0])
                                     for ff in time])
                else:
                    raise PyLCInputError(
                        'Not valid time format. Available formats: JD_UTC, MJD_UTC, BJD_UTC, BJD_TDB, BJD_TT, '
                        'HJD_UTC, HJD_BJD, HJD_TT')

            except:
                raise PyLCInputError('Not valid input for time')

    def convert_to_jd(self, time, time_format):

        if isinstance(time, float):

            if time_format in ['BJD_TDB', 'BJD_TT']:
                return curve_fit(self._bjd_tdb, [0], [time], p0=[time])[0][0]
            elif time_format == 'JD_UTC':
                return time
            elif time_format == 'MJD_UTC':
                return time + 2400000.5
            elif time_format == 'BJD_UTC':
                return curve_fit(self._bjd_utc, [0], [time], p0=[time])[0][0]
            elif time_format in ['HJD_TDB', 'HJD_TT']:
                return curve_fit(self._hjd_tdb, [0], [time], p0=[time])[0][0]
            elif time_format == 'HJD_UTC':
                return curve_fit(self._hjd_utc, [0], [time], p0=[time])[0][0]
            else:
                raise PyLCInputError(
                    'Not valid time format. Available formats: JD_UTC, MJD_UTC, BJD_UTC, BJD_TDB, BJD_TT, '
                    'HJD_UTC, HJD_BJD, HJD_TT')

        else:
            try:
                time = np.array(time, dtype=float)

                if time_format in ['BJD_TDB', 'BJD_TT']:
                    return np.array([curve_fit(self._bjd_tdb, [0], [ff], p0=[time])[0][0] for ff in time])
                elif time_format == 'JD_UTC':
                    return time
                elif time_format == 'MJD_UTC':
                    return time + 2400000.5
                elif time_format == 'BJD_UTC':
                    return np.array([curve_fit(self._bjd_utc, [0], [ff], p0=[time])[0][0] for ff in time])
                elif time_format in ['HJD_TDB', 'HJD_TT']:
                    return np.array([curve_fit(self._hjd_tdb, [0], [ff], p0=[time])[0][0] for ff in time])
                elif time_format == 'HJD_UTC':
                    return np.array([curve_fit(self._hjd_utc, [0], [ff], p0=[time])[0][0] for ff in time])
                else:
                    raise PyLCInputError(
                        'Not valid time format. Available formats: JD_UTC, MJD_UTC, BJD_UTC, BJD_TDB, BJD_TT, '
                        'HJD_UTC, HJD_BJD, HJD_TT')

            except:
                raise PyLCInputError('Not valid input for time')

    def convert_to_mjd(self, time, time_format):

        return self.convert_to_jd(time, time_format) - 2400000.5

    def _hjd_utc(self, x, jd):

        ssb_ra, ssb_dec, ssb_d, ssb_dt = plc_data.heliocentre(jd)

        a = ssb_d / 60.0 / 24.0
        b = self.dec.sin() * np.sin(ssb_dec)
        c = self.dec.cos() * np.cos(ssb_dec) * np.cos(self.ra.rad() - ssb_ra)

        return jd - a * (b + c)

    def _hjd_tdb(self, x, jd):

        ssb_ra, ssb_dec, ssb_d, ssb_dt = plc_data.heliocentre(jd)

        a = ssb_d / 60.0 / 24.0
        b = self.dec.sin() * np.sin(ssb_dec)
        c = self.dec.cos() * np.cos(ssb_dec) * np.cos(self.ra.rad() - ssb_ra)

        return jd - a * (b + c) + ssb_dt / 60.0 / 60.0 / 24.0

    def _bjd_utc(self, x, jd):

        ssb_ra, ssb_dec, ssb_d, ssb_dt = plc_data.barycentre(jd)

        a = ssb_d / 60.0 / 24.0
        b = self.dec.sin() * np.sin(ssb_dec)
        c = self.dec.cos() * np.cos(ssb_dec) * np.cos(self.ra.rad() - ssb_ra)

        return jd - a * (b + c)

    def _bjd_tdb(self, x, jd):

        ssb_ra, ssb_dec, ssb_d, ssb_dt = plc_data.barycentre(jd)

        a = ssb_d / 60.0 / 24.0
        b = self.dec.sin() * np.sin(ssb_dec)
        c = self.dec.cos() * np.cos(ssb_dec) * np.cos(self.ra.rad() - ssb_ra)

        return jd - a * (b + c) + ssb_dt / 60.0 / 60.0 / 24.0

    def distance_on_sphere(self, other):

        _request_target(other)

        return arccos(self.dec.sin() * other.dec.sin() + self.dec.cos() * other.dec.cos() * (self.ra - other.ra).cos())

    def __str__(self):
        return 'plc.FixedTarget(RA(hms)/DEC(dms): {0})'.format(self.coord)

    def __repr__(self):
        return self.__str__()

def _is_target(item):
    if isinstance(item, FixedTarget):
        return True
    else:
        return False

def _request_target(item):
    if _is_target(item):
        pass
    else:
        raise PyLCInputError('A plc.Target object is required (plc.FixedTarget)')