#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Thu Feb  7 16:55:54 2019
Module for the SitePair class.
@author: parrenif
"""

import os
import sys
import numpy as np
import matplotlib.pyplot as mpl
from matplotlib.backends.backend_pdf import PdfPages
from scipy.linalg import lu_factor, lu_solve
from scipy.linalg import cholesky
import pccfg

class SitePair(object):
    """Class for a pair of sites."""

    def __init__(self, site1, site2):
        self.site1 = site1
        self.site2 = site2
        self.label = self.site1.label+'-'+self.site2.label

        self.age_age_label = self.site1.age_label+self.site2.age_label
        if len(self.age_age_label)>0:
            self.age_age_label = self.age_age_label + '_'
        self.age_age2_label = self.site1.age_label+self.site2.age2_label + '_'
        self.age2_age_label = self.site1.age2_label+self.site2.age_label + '_'
        self.age2_age2_label = self.site1.age2_label+self.site2.age2_label + '_'

#TODO: allow to have either dlabel1+'-'dlabel2 or dlbel2+'-'dlabel1 as directory
        filename = pccfg.datadir+self.site1.label+'-'+self.site2.label+\
            '/'+self.age_age_label+'synchro_horizons.txt'
        if not os.path.isfile(filename):
            filename = pccfg.datadir+self.site1.label+'-'+self.site2.label+'/ice_depth.txt'
        if os.path.isfile(filename) and open(filename).read():
            readarray = np.loadtxt(filename)
            if np.size(readarray) == np.shape(readarray)[0]:
                readarray.resize(1, np.size(readarray))
            self.iceicehorizons_depth1 = readarray[:, 0]
            self.iceicehorizons_depth2 = readarray[:, 1]
            self.iceicehorizons_sigma = readarray[:, 2]
        else:
            self.iceicehorizons_depth1 = np.array([])
            self.iceicehorizons_depth2 = np.array([])
            self.iceicehorizons_sigma = np.array([])
        self.iceicehorizons_correlation = np.diag(np.ones(np.size(self.iceicehorizons_depth1)))

        if self.site1.archive == 'icecore' and self.site2.archive == 'icecore':
            filename = pccfg.datadir+self.site1.label+'-'+self.site2.label+\
                       '/'+self.age2_age2_label+'synchro_horizons.txt'
            if not os.path.isfile(filename):
                filename = pccfg.datadir+self.site1.label+'-'+self.site2.label+'/air_depth.txt'
            if os.path.isfile(filename) and open(filename).read():
                readarray = np.loadtxt(filename)
                if np.size(readarray) == np.shape(readarray)[0]:
                    readarray.resize(1, np.size(readarray))
                self.airairhorizons_depth1 = readarray[:, 0]
                self.airairhorizons_depth2 = readarray[:, 1]
                self.airairhorizons_sigma = readarray[:, 2]
            else:
                self.airairhorizons_depth1 = np.array([])
                self.airairhorizons_depth2 = np.array([])
                self.airairhorizons_sigma = np.array([])
            self.airairhorizons_correlation = np.diag(np.ones(np.size(self.airairhorizons_depth1)))

        if self.site2.archive == 'icecore':
            filename = pccfg.datadir+self.site1.label+'-'+\
                            self.site2.label+'/'+self.age_age2_label+'synchro_horizons.txt'
            if not os.path.isfile(filename):
                filename = pccfg.datadir+self.site1.label+'-'+self.site2.label+'/iceair_depth.txt'
            if os.path.isfile(filename) and open(filename).read():
                readarray = np.loadtxt(filename)
                if np.size(readarray) == np.shape(readarray)[0]:
                    readarray.resize(1, np.size(readarray))
                self.iceairhorizons_depth1 = readarray[:, 0]
                self.iceairhorizons_depth2 = readarray[:, 1]
                self.iceairhorizons_sigma = readarray[:, 2]
            else:
                self.iceairhorizons_depth1 = np.array([])
                self.iceairhorizons_depth2 = np.array([])
                self.iceairhorizons_sigma = np.array([])
            self.iceairhorizons_correlation = np.diag(np.ones(np.size(self.iceairhorizons_depth1)))

        if self.site1.archive == 'icecore':
            filename = pccfg.datadir+self.site1.label+'-'+\
                        self.site2.label+'/'+self.age2_age_label+'synchro_horizons.txt'
            if not os.path.isfile(filename):
                filename = pccfg.datadir+self.site1.label+'-'+self.site2.label+'/airice_depth.txt'
            if os.path.isfile(filename) and open(filename).read():
                readarray = np.loadtxt(filename)
                if np.size(readarray) == np.shape(readarray)[0]:
                    readarray.resize(1, np.size(readarray))
                self.airicehorizons_depth1 = readarray[:, 0]
                self.airicehorizons_depth2 = readarray[:, 1]
                self.airicehorizons_sigma = readarray[:, 2]
            else:
                self.airicehorizons_depth1 = np.array([])
                self.airicehorizons_depth2 = np.array([])
                self.airicehorizons_sigma = np.array([])
            self.airicehorizons_correlation = np.diag(np.ones(np.size(self.airicehorizons_depth1)))

        filename1 = pccfg.datadir+'/parameters_covariance_observations_all_site_pairs.py'
        filename2 = pccfg.datadir+'/parameters-CovarianceObservations-AllDrillings.py'
        if os.path.isfile(filename1):
            exec(open(filename1).read())
        elif os.path.isfile(filename2):
            exec(open(filename2).read())
        filename3 = pccfg.datadir+self.label+'/parameters_covariance_observations.py'
        filename4 = pccfg.datadir+self.label+'/parameters-CovarianceObservations.py'
        if os.path.isfile(filename3):
            exec(open(filename3).read())
        elif os.path.isfile(filename4):
            exec(open(filename4).read())

        if ((os.path.isfile(filename1) or os.path.isfile(filename2) or os.path.isfile(filename3)\
            or os.path.isfile(filename4)) and (pccfg.jacobian=='analytical' or \
            pccfg.jacobian=='semi_adjoint' or pccfg.jacobian=='adjoint')):
            print('Covariance on observations not implemented for analytical Jacobian. Exiting.')
            sys.exit()

        if np.any(self.iceicehorizons_correlation != \
                  np.diag(np.ones(np.size(self.iceicehorizons_depth1)))):
            self.iceicehorizons_correlation_bool = True
            self.iceicehorizons_chol = cholesky(self.iceicehorizons_correlation)
            self.iceicehorizons_lu_piv = lu_factor(self.iceicehorizons_chol)
        else:
            self.iceicehorizons_correlation_bool = False

        if self.site1.archive == 'icecore' and self.site2.archive == 'icecore':
            if np.any(self.airairhorizons_correlation != \
                  np.diag(np.ones(np.size(self.airairhorizons_depth1)))):
                self.airairhorizons_correlation_bool = True
                self.airairhorizons_chol = cholesky(self.airairhorizons_correlation)
                self.airairhorizons_lu_piv = lu_factor(self.airairhorizons_chol)
            else:
                self.airairhorizons_correlation_bool = False

        if self.site2.archive == 'icecore':
            if np.any(self.iceairhorizons_correlation != \
                  np.diag(np.ones(np.size(self.iceairhorizons_depth1)))):
                self.iceairhorizons_correlation_bool = True
                self.iceairhorizons_chol = cholesky(self.iceairhorizons_correlation)
                self.iceairhorizons_lu_piv = lu_factor(self.iceairhorizons_chol)
            else:
                self.iceairhorizons_correlation_bool = False

        if self.site1.archive == 'icecore':
            if np.any(self.airicehorizons_correlation != \
                  np.diag(np.ones(np.size(self.airicehorizons_depth1)))):
                self.airicehorizons_correlation_bool = True
                self.airicehorizons_chol = cholesky(self.airicehorizons_correlation)
                self.airicehorizons_lu_piv = lu_factor(self.airicehorizons_chol)
            else:
                self.airicehorizons_correlation_bool = False

    def residuals(self):
        """Calculate the residual terms of a pair of sites."""

        if np.size(self.iceicehorizons_depth1) > 0:
            resi_iceice = (self.site1.fct_age(self.iceicehorizons_depth1)-\
                           self.site2.fct_age(self.iceicehorizons_depth2))/self.iceicehorizons_sigma
            if self.iceicehorizons_correlation_bool:
                resi_iceice = lu_solve(self.iceicehorizons_lu_piv, resi_iceice)
            resi = [resi_iceice]
        else:
            resi = [np.array([])]

        if self.site1.archive == 'icecore' and self.site2.archive == 'icecore' and \
            np.size(self.airairhorizons_depth1) > 0:
            resi_airair = (self.site1.fct_airage(self.airairhorizons_depth1)-\
                          self.site2.fct_airage(self.airairhorizons_depth2))/\
                          self.airairhorizons_sigma
            if self.airairhorizons_correlation_bool:
                resi_airair = lu_solve(self.airairhorizons_lu_piv, resi_airair)
            resi.append(resi_airair)

        if self.site2.archive == 'icecore' and np.size(self.iceairhorizons_depth1) > 0:
            resi_iceair = (self.site1.fct_age(self.iceairhorizons_depth1)-\
                          self.site2.fct_airage(self.iceairhorizons_depth2))/\
                          self.iceairhorizons_sigma
            if self.iceairhorizons_correlation_bool:
                resi_iceair = lu_solve(self.iceairhorizons_lu_piv, resi_iceair)
            resi.append(resi_iceair)

        if self.site1.archive == 'icecore' and np.size(self.airicehorizons_depth1) > 0:
            resi_airice = (self.site1.fct_airage(self.airicehorizons_depth1)-\
                           self.site2.fct_age(self.airicehorizons_depth2))/self.airicehorizons_sigma
            if self.airicehorizons_correlation_bool:
                resi_airice = lu_solve(self.airicehorizons_lu_piv, resi_airice)
            resi.append(resi_airice)

        return np.concatenate(resi)

    def residuals_jacobian1(self):
#        if np.size(self.iceicehorizons_depth1) > 0:
#            print(np.shape(self.site1.fct_age_jac(self.iceicehorizons_depth1)),
#                  np.shape(self.site2.fct_age(self.iceicehorizons_depth2)),
#                  np.shape(self.iceicehorizons_sigma)
        resi_iceice = self.site1.fct_age_jac(self.iceicehorizons_depth1)/self.iceicehorizons_sigma
        resi = [resi_iceice]
        if self.site1.archive == 'icecore' and self.site2.archive == 'icecore':
            resi_airair = self.site1.fct_airage_jac(self.airairhorizons_depth1)/\
                            self.airairhorizons_sigma
            resi.append(resi_airair)
        if self.site2.archive == 'icecore':
            resi_iceair = self.site1.fct_age_jac(self.iceairhorizons_depth1)/\
                            self.iceairhorizons_sigma
            resi.append(resi_iceair)
        if self.site1.archive == 'icecore':
            resi_airice = self.site1.fct_airage_jac(self.airicehorizons_depth1)/\
                            self.airicehorizons_sigma
            resi.append(resi_airice)
#       else:
#            resi = [np.array([])]
        return np.concatenate(resi, axis=1)

    def residuals_jacobian2(self):
#        if np.size(self.iceicehorizons_depth1) > 0:
        resi_iceice = -self.site2.fct_age_jac(self.iceicehorizons_depth2)/self.iceicehorizons_sigma
        resi = [resi_iceice]
        if self.site1.archive == 'icecore' and self.site2.archive == 'icecore':
            resi_airair = -self.site2.fct_airage_jac(self.airairhorizons_depth2)/\
                            self.airairhorizons_sigma
            resi.append(resi_airair)
        if self.site2.archive == 'icecore':
            resi_iceair = -self.site2.fct_airage_jac(self.iceairhorizons_depth2)/\
                            self.iceairhorizons_sigma
            resi.append(resi_iceair)
        if self.site1.archive == 'icecore':
            resi_airice = -self.site2.fct_age_jac(self.airicehorizons_depth2)/\
                            self.airicehorizons_sigma
            resi.append(resi_airice)
#        else:
#            resi = [np.array([])]
        return np.concatenate(resi, axis=1)

    def residuals_delta(self):
        resi_iceice = (self.site1.fct_age_delta(self.iceicehorizons_depth1)-\
            self.site2.fct_age_delta(self.iceicehorizons_depth2))/self.iceicehorizons_sigma
        return resi_iceice

    def figures(self):

        """Build the figures related to a pair of sites."""
        if np.size(self.iceicehorizons_depth1)>0:
            fig, ax = mpl.subplots()
            mpl.xlabel(self.site1.label+' '+self.site1.age_labelsp+'age ('+pccfg.age_unit+' '+pccfg.age_unit_ref+')')
            mpl.ylabel(self.site2.label+' '+self.site2.age_labelsp+'age ('+pccfg.age_unit+' '+pccfg.age_unit_ref+')')
            if np.size(self.iceicehorizons_depth1) > 0:
                if pccfg.show_initial:
                    mpl.plot(self.site1.fct_age_init(self.iceicehorizons_depth1),
                                 self.site2.fct_age_init(self.iceicehorizons_depth2),
                                 color=pccfg.color_init, linestyle='', marker='o', markersize=2,
                                 label="Initial")
                mpl.plot(self.site1.fct_age_model(self.iceicehorizons_depth1),
                             self.site2.fct_age_model(self.iceicehorizons_depth2),
                             color=pccfg.color_mod, linestyle='', marker='o', markersize=2,
                             label="Prior")
                mpl.errorbar(self.site1.fct_age(self.iceicehorizons_depth1),
                             self.site2.fct_age(self.iceicehorizons_depth2), color=pccfg.color_opt,
                             ecolor=pccfg.color_ci,
                             xerr=np.zeros(np.size(self.iceicehorizons_depth1)),
                             linestyle='', marker='o', markersize=2,
                             label="Posterior $\pm\sigma$")
                xstart = self.site1.fct_age(self.iceicehorizons_depth1)-self.iceicehorizons_sigma/2
                ystart = self.site2.fct_age(self.iceicehorizons_depth2)+self.iceicehorizons_sigma/2
                for i in range(np.size(self.iceicehorizons_depth1)):
                    mpl.arrow(xstart[i], ystart[i], self.iceicehorizons_sigma[i],
                              -self.iceicehorizons_sigma[i], color=pccfg.color_ci,
                              width=0.0, head_length=0.0, head_width=0.0)
            x_low, x_up, y_low, y_up = mpl.axis()
#            x_low = self.site1.age_top
#            y_low = self.site2.age_top
#            mpl.axis((x_low, x_up, y_low, y_up))
            rangefig = np.array([min(x_low, y_low), max(x_up, y_up)])
            mpl.plot(rangefig, rangefig, color=pccfg.color_obs, label='Perfect agreement', zorder=0)
            mpl.legend(loc="best")
            ax.set_aspect('equal')
            printed_page = PdfPages(pccfg.datadir+self.label+'/'+self.age_age_label+'synchro.pdf')
            printed_page.savefig(fig)
            printed_page.close()
            if not pccfg.show_figures:
                mpl.close()

        if self.site1.archive == 'icecore' and self.site2.archive == 'icecore':
            if np.size(self.airairhorizons_depth1)>0:
                fig, ax = mpl.subplots()
                mpl.xlabel(self.site1.label+' '+self.site1.age2_labelsp+'age ('+pccfg.age_unit+' '+
                           pccfg.age_unit_ref+')')
                mpl.ylabel(self.site2.label+' '+self.site2.age2_labelsp+'age ('+pccfg.age_unit+' '+
                           pccfg.age_unit_ref+')')
                if np.size(self.airairhorizons_depth1) > 0:
                    if pccfg.show_initial:
                        mpl.plot(self.site1.fct_airage_init(self.airairhorizons_depth1),
                                     self.site2.fct_airage_init(self.airairhorizons_depth2),
                                     color=pccfg.color_init,
                                     linestyle='',
                                     marker='o', markersize=2, label="Initial")
                    mpl.plot(self.site1.fct_airage_model(self.airairhorizons_depth1),
                                 self.site2.fct_airage_model(self.airairhorizons_depth2),
                                 color=pccfg.color_mod,
                                 linestyle='', marker='o', markersize=2,
                                 label="Prior")
                    mpl.errorbar(self.site1.fct_airage(self.airairhorizons_depth1),
                                 self.site2.fct_airage(self.airairhorizons_depth2),
                                 color=pccfg.color_opt, ecolor=pccfg.color_ci,
                                 xerr=np.zeros_like(self.airairhorizons_sigma),
                                 linestyle='', marker='o', markersize=2,
                                 label="Posterior $\pm\sigma$")
                    xstart = self.site1.fct_airage(self.airairhorizons_depth1)-\
                                 self.airairhorizons_sigma/2
                    ystart = self.site2.fct_airage(self.airairhorizons_depth2)+\
                                 self.airairhorizons_sigma/2
                    for i in range(np.size(self.airairhorizons_depth1)):
                        mpl.arrow(xstart[i], ystart[i], self.airairhorizons_sigma[i],
                                  -self.airairhorizons_sigma[i], color=pccfg.color_ci,
                                  width=0.0, head_length=0.0, head_width=0.0)
                x_low, x_up, y_low, y_up = mpl.axis()
#                x_low = self.site1.age_top
#                y_low = self.site2.age_top
#                mpl.axis((x_low, x_up, y_low, y_up))
                rangefig = np.array([min(x_low, y_low), max(x_up, y_up)])
                mpl.plot(rangefig, rangefig, color=pccfg.color_obs, label='Perfect agreement',
                         zorder=0)
                mpl.legend(loc="best")
                ax.set_aspect('equal')
                printed_page = PdfPages(pccfg.datadir+self.label+'/'+self.age2_age2_label+
                                        'synchro.pdf')
                printed_page.savefig(fig)
                printed_page.close()
                if not pccfg.show_figures:
                    mpl.close()

        if self.site2.archive == 'icecore':
            if np.size(self.iceairhorizons_depth1)>0:
                fig, ax = mpl.subplots()
                mpl.xlabel(self.site1.label+' '+self.site1.age_labelsp+'age ('+pccfg.age_unit+' '+
                           pccfg.age_unit_ref+')')
                mpl.ylabel(self.site2.label+' '+self.site2.age2_labelsp+'age ('+pccfg.age_unit+' '+
                           pccfg.age_unit_ref+')')
                if np.size(self.iceairhorizons_depth1) > 0:
                    if pccfg.show_initial:
                        mpl.plot(self.site1.fct_age_init(self.iceairhorizons_depth1),
                                     self.site2.fct_airage_init(self.iceairhorizons_depth2),
                                     color=pccfg.color_init,
                                     linestyle='',
                                     marker='o', markersize=2, label="Initial")
                    mpl.plot(self.site1.fct_age_model(self.iceairhorizons_depth1),
                                 self.site2.fct_airage_model(self.iceairhorizons_depth2),
                                 color=pccfg.color_mod,
                                 linestyle='', marker='o', markersize=2,
                                 label="Prior")
                    mpl.errorbar(self.site1.fct_age(self.iceairhorizons_depth1),
                                 self.site2.fct_airage(self.iceairhorizons_depth2),
                                 color=pccfg.color_opt, ecolor=pccfg.color_ci,
                                 xerr=np.zeros_like(self.iceairhorizons_sigma),
                                 linestyle='', marker='o', markersize=2,
                                 label="Posterior $\pm\sigma$")
                    xstart = self.site1.fct_age(self.iceairhorizons_depth1)-\
                                 self.iceairhorizons_sigma/2
                    ystart = self.site2.fct_airage(self.iceairhorizons_depth2)+\
                                 self.iceairhorizons_sigma/2
                    for i in range(np.size(self.iceairhorizons_depth1)):
                        mpl.arrow(xstart[i], ystart[i], self.iceairhorizons_sigma[i],
                                  -self.iceairhorizons_sigma[i], color=pccfg.color_ci,
                                  width=0.0, head_length=0.0, head_width=0.0)
                x_low, x_up, y_low, y_up = mpl.axis()
#                x_low = self.site1.age_top
#                y_low = self.site2.age_top
#                mpl.axis((x_low, x_up, y_low, y_up))
                rangefig = np.array([min(x_low, y_low), max(x_up, y_up)])
                mpl.plot(rangefig, rangefig, color=pccfg.color_obs, label='Perfect agreement',
                         zorder=0)
                mpl.legend(loc="best")
                ax.set_aspect('equal')
                printed_page = PdfPages(pccfg.datadir+self.label+'/'+self.age_age2_label+
                                        'synchro.pdf')
                printed_page.savefig(fig)
                printed_page.close()
                if not pccfg.show_figures:
                    mpl.close()

        if self.site1.archive == 'icecore':
            if np.size(self.airicehorizons_depth1)>0:
                fig, ax = mpl.subplots()
                mpl.xlabel(self.site1.label+' '+self.site1.age2_labelsp+'age ('+pccfg.age_unit+' '+
                           pccfg.age_unit_ref+')')
                mpl.ylabel(self.site2.label+' '+self.site2.age_labelsp+'age ('+pccfg.age_unit+' '+
                           pccfg.age_unit_ref+')')
                if np.size(self.airicehorizons_depth1) > 0:
                    if pccfg.show_initial:
                        mpl.plot(self.site1.fct_airage_init(self.airicehorizons_depth1),
                                     self.site2.fct_age_init(self.airicehorizons_depth2),
                                     color=pccfg.color_init,
                                     linestyle='', marker='o', markersize=2, label="Initial")
                    mpl.plot(self.site1.fct_airage_model(self.airicehorizons_depth1),
                                 self.site2.fct_age_model(self.airicehorizons_depth2),
                                 color=pccfg.color_mod,
                                 linestyle='', marker='o', markersize=2,
                                 label="Prior")
                    mpl.errorbar(self.site1.fct_airage(self.airicehorizons_depth1),
                                 self.site2.fct_age(self.airicehorizons_depth2),
                                 color=pccfg.color_opt, ecolor=pccfg.color_ci,
                                 xerr=np.zeros_like(self.airicehorizons_sigma),
                                 linestyle='', marker='o', markersize=2,
                                 label="Posterior $\pm\sigma$")
                    xstart = self.site1.fct_airage(self.airicehorizons_depth1)-\
                                 self.airicehorizons_sigma/2
                    ystart = self.site2.fct_age(self.airicehorizons_depth2)+\
                                 self.airicehorizons_sigma/2
                    for i in range(np.size(self.airicehorizons_depth1)):
                        mpl.arrow(xstart[i], ystart[i], self.airicehorizons_sigma[i],
                                  -self.airicehorizons_sigma[i], color=pccfg.color_ci,
                                  width=0.0, head_length=0.0, head_width=0.0)
                x_low, x_up, y_low, y_up = mpl.axis()
#                x_low = self.site1.age_top
#                y_low = self.site2.age_top
#                mpl.axis((x_low, x_up, y_low, y_up))
                rangefig = np.array([min(x_low, y_low), max(x_up, y_up)])
                mpl.plot(rangefig, rangefig, color=pccfg.color_obs, label='Perfect agreement')
                mpl.legend(loc="best")
                ax.set_aspect('equal')
                printed_page = PdfPages(pccfg.datadir+self.label+'/'+self.age2_age_label+
                                        'synchro.pdf')
                printed_page.savefig(fig)
                printed_page.close()
                if not pccfg.show_figures:
                    mpl.close()


# Licensed with the 3-clause BSD license.  See LICENSE for details.
from itertools import groupby
import numpy as np
import astropy.units as u
from astropy.time import Time
from astropy.table import vstack
from sbpy.data import Ephem, Orbit, Names, QueryError
from sbpy.data.names import TargetNameParseError
from . import util

STEP_LIMIT = 300

def generate(desg, location, epochs, source='jpl', orbit=None, cache=False):
    """Generate ephemeris at specific epochs from external source.

    Parameters
    ----------
    desg : string
        Object designation.

    location : string
        Observer location.

    epochs : array-like, Time, or dict
        Compute ephemeris at these epochs.  Arrays must be
        floats (for Julian date) or else parsable by
        `~astropy.time.Time`.  Dictionaries are passed to the
        ephemeris source as is.  Set ``step=None`` for an
        adaptable time step that mitigates interpolation errors.

    source : string, optional
        Source to use: 'mpc', 'jpl', or 'oorb'.  'oorb' requires
        ``orbit`` parameter.

    orbit : `~sbpy.data.Orbit`, optional
        Orbital elements for ``source=oorb``.

    cache : bool, optional
        Use cached ephemerides; primarily for testing.

    Returns
    -------
    eph : `~sbpy.data.ephem.Ephem`
        Ephemeris.

    """

    if source not in ['mpc', 'jpl', 'oorb']:
        raise ValueError('Source must be "mpc", "jpl", or "oorb".')

    _epochs = _format_epochs(epochs)
    if isinstance(_epochs, dict):
        if epochs.get('step') is None:
            eph = _get_adaptable_steps(desg, location, _epochs,
                                       source=source, orbit=orbit,
                                       cache=cache)
    else:
        eph = _get_fixed_steps(desg, location, _epochs, source=source,
                               orbit=orbit, cache=cache)
    return eph

def generate_orbit(desg, epochs, cache=False):
    """Generate orbital parameters from JPL Horizons at specific epochs.

    Parameters
    ----------
    desg: string
        Object designation.

    epochs: array-like or dict
        Compute orbital elements at these epochs.  For arrays,
        must be floats(for Julian date) or else parsable by
        `~astropy.time.Time`.

    cache: bool, optional
        Use cached ephemerides; primarily for testing.

    Returns
    -------
    orb: `~sbpy.data.Orbit`
        Orbital elements.

    """

    _epochs = _format_epochs(epochs)
    if not isinstance(_epochs, dict):
        if len(_epochs) > STEP_LIMIT:
            orb = None
            N = np.ceil(len(_epochs) / STEP_LIMIT)
            for e in np.array_split(_epochs, N):
                _orb = generate_orbit(desg, e, cache=cache)
                if orb:
                    orb.add_rows(_orb)
                else:
                    orb = _orb
            return orb
        else:
            pass
            # ... and continue

    kwargs = dict(epochs=_epochs, cache=cache)
    kwargs['id_type'] = 'smallbody'
    comet = False
    try:
        Names.parse_comet(desg)
        comet = True
        kwargs['id_type'] = 'designation'
        cap_limit = closest_apparition_limit(_epochs)
        kwargs.update(closest_apparition=cap_limit,
                      no_fragments=True)
    except TargetNameParseError:
        pass

    try:
        orb = Orbit.from_horizons(desg, **kwargs)
    except QueryError:
        # Dual-listed objects should be queried without CAP/NOFRAG.  If
        # this was a comet query, try again.
        if comet:
            del kwargs['closest_apparition'], kwargs['no_fragments']

            # fix for sbpy bug #xxx, remove once v0.2.2 is released
            kwargs['epochs'] = _format_epochs(epochs)

            orb = Orbit.from_horizons(desg, **kwargs)
        else:
            raise

    return orb

def closest_apparition_limit(epochs):
    if isinstance(epochs, dict):
        start = epochs['start'].jd
    else:
        start = epochs[0].jd
    return '<{}'.format(start)

def _format_epochs(epochs):
    if isinstance(epochs, dict):
        start, stop = util.epochs_to_time((epochs['start'], epochs['stop']))
        step = (None if epochs.get('step') is None
                else u.Quantity(epochs.get('step')))
        e = {
            'start': start,
            'stop': stop,
            'step': step
        }
    else:
        e = util.epochs_to_time(epochs)
        if len(epochs) > 1:
            d = np.diff(e.jd)
            if any(d <= 0):
                raise ValueError('Epoch dates must be increasing and unique.')

    return e

def _get_fixed_steps(desg, location, epochs, source='jpl', orbit=None,
                     cache=False):
    if not isinstance(epochs, dict):
        # list of specific dates, divide into chunks, as needed
        if len(epochs) > STEP_LIMIT:
            eph = None
            N = np.ceil(len(epochs) / STEP_LIMIT)
            for e in np.array_split(epochs, N):
                _eph = _get_fixed_steps(desg, location, list(e),
                                        source=source, cache=cache)
                if eph:
                    eph.add_rows(_eph)
                else:
                    eph = _eph
            return eph
        else:
            pass
            # and proceed...

    if source == 'mpc':
        eph = Ephem.from_mpc(desg, epochs=epochs,
                             location=location,
                             proper_motion='sky',
                             proper_motion_unit='rad/s',
                             cache=cache)

        z = np.zeros(len(eph))
        if 'Uncertainty 3sig' not in eph.table.colnames:
            eph.table.add_column(u.Quantity(z, 'arcsec'),
                                 name='SMAA_3sigma')
            eph.table.add_column(u.Quantity(z, 'arcsec'),
                                 name='SMIA_3sigma')
            eph.table.add_column(u.Quantity(z, 'rad'),
                                 name='Theta_3sigma')
        else:
            # MPC's ephemeris uncertainty is a line, rather than
            # an ellipse
            eph.table.add_column(u.Quantity(z, 'arcsec'),
                                 name='SMIA_3sigma')
            eph.table.add_column(eph['Uncertainty 3sig'],
                                 name='SMAA_3sigma')
            eph.table.add_column(eph['Unc. P.A.'],
                                 name='Theta_3sigma')

        # convert date from Time to Julian date, this helps with table
        # manipulation since Time columns cannot be appended to.
        eph['jd'] = eph['date'].jd
    elif source == 'jpl':
        # column 7 (sidereal time) isn't needed, but adding to avoid sbpy-0.2.1 crash
        kwargs = dict(
            epochs=epochs,
            location=location,
            quantities='1,3,8,9,7,19,20,23,24,27,36,37',
            cache=cache
        )
        kwargs['id_type'] = 'smallbody'
        comet = False
        try:
            Names.parse_comet(desg)
            comet = True
            kwargs['id_type'] = 'designation'
            cap_limit = closest_apparition_limit(epochs)
            kwargs.update(
                closest_apparition=cap_limit,
                no_fragments=True
            )
        except TargetNameParseError:
            pass

        try:
            eph = Ephem.from_horizons(desg, **kwargs)
        except QueryError:
            # Dual-listed objects should be queried without CAP/NOFRAG.  If
            # this was a comet query, try again.
            if comet:
                del kwargs['closest_apparition'], kwargs['no_fragments']

                # fix for sbpy bug #xxx, remove once v0.2.2 is released
                kwargs['epochs'] = _format_epochs(epochs)

                eph = Ephem.from_horizons(desg, **kwargs)
            else:
                raise

        # Horizon's theta is measured N of E.  We want E of N.
        eph['Theta_3sigma'] = (np.pi / 2) * u.rad - eph['Theta_3sigma']

        # create a plain Julian date column, helps with table
        # manipulation in _get_adaptable_steps
        eph['jd'] = eph['epoch'].jd
    elif source == 'oorb':
        eph = Ephem.from_oo(orbit, epochs=epochs, location=location)
        # no uncertainties from oorb
        z = np.zeros(len(eph))
        eph.table.add_column(u.Quantity(z, 'arcsec'),
                             name='SMAA_3sigma')
        eph.table.add_column(u.Quantity(z, 'arcsec'),
                             name='SMIA_3sigma')
        eph.table.add_column(u.Quantity(z, 'rad'),
                             name='Theta_3sigma')
        eph['jd'] = eph['epoch'].jd  # same note as for jpl

    return eph

def _get_adaptable_steps(desg, location, epochs, source='jpl', orbit=None,
                         cache=False):
    """Ephemerides with adaptable time step.

    Based on ZChecker analysis, Oct 2018, the interpolation error is
    ~2" / delta for 6h time step.

    """

    # daily ephemeris for delta > 1
    epochs['step'] = 1 * u.day
    eph = _get_fixed_steps(desg, location, epochs, source=source,
                           cache=cache)

    for limit, substep in ((1 * u.au, 4 * u.hr), (0.25 * u.au, 1 * u.hr)):
        updated = []
        groups = groupby(zip(eph['delta'], eph['jd']),
                         lambda e: e[0] < limit)
        for inside, epochs in groups:
            if not inside:
                continue

            jd = [e[1] for e in epochs]
            if len(jd) > 1:
                sub_epochs = _format_epochs({
                    'start': jd[0],
                    'stop': jd[-1],
                    'step': substep
                })
                update = _get_fixed_steps(desg, location, sub_epochs,
                                          source=source, cache=cache)
                updated.append(update.table)

        eph = Ephem.from_table(vstack([eph.table] + updated))

        # remove duplicates
        eph.table.sort('jd')
        d = np.diff(eph['jd'])
        duplicates = np.flatnonzero(np.isclose(d, 0)) + 1
        eph.table.remove_rows(duplicates)

    return eph

import numpy,import astropy,from astropy


#!/usr/bin/env python
# Copyright 2014-2019 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
QM part interface
'''

import numpy
import pyscf
from pyscf import lib
from pyscf import gto
from pyscf import df
from pyscf import scf
from pyscf import mcscf
from pyscf import grad
from pyscf.lib import logger

def mm_charge(scf_method, coords, charges, unit=None):
    '''Modify the QM method using the (non-relativistic) potential generated
    by MM charges. Note the static Coulomb interactions of the background
    charges are not included in the total energy.

    Args:
        scf_method : a HF or DFT object

        coords : 2D array, shape (N,3)
            MM particle coordinates
        charges : 1D array
            MM particle charges
    Kwargs:
        unit : str
            Bohr, AU, Ang (case insensitive). Default is the same to mol.unit

    Returns:
        Same method object as the input scf_method with modified 1e Hamiltonian

    Note:
        1. if MM charge and X2C correction are used together, function mm_charge
        needs to be applied after X2C decoration (.x2c method), eg
        mf = mm_charge(scf.RHF(mol).x2c()), [(0.5,0.6,0.8)], [-0.5]).
        2. Once mm_charge function is applied on the SCF object, it
        affects all the post-HF calculations eg MP2, CCSD, MCSCF etc

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = mm_charge(dft.RKS(mol), [(0.5,0.6,0.8)], [-0.3])
    >>> mf.kernel()
    -101.940495711284
    '''
    assert(isinstance(scf_method, scf.hf.SCF) or
           isinstance(scf_method, mcscf.casci.CASCI))

    if unit is None:
        unit = scf_method.mol.unit
    if unit.startswith(('B','b','au','AU')):
        coords = numpy.asarray(coords, order='C')
    elif unit.startswith(('A','a')):
        coords = numpy.asarray(coords, order='C') / lib.parameters.BOHR
    else:
        coords = numpy.asarray(coords, order='C') / unit
    charges = numpy.asarray(charges)
    method_class = scf_method.__class__

    class QMMM(method_class, _QMMM):
        def __init__(self):
            self.__dict__.update(scf_method.__dict__)

        def dump_flags(self, verbose=None):
            method_class.dump_flags(self, verbose)
            logger.info(self, '** Add background charges for %s **',
                        method_class)
            if self.verbose >= logger.DEBUG:
                logger.debug(self, 'Charge      Location')
                for i, z in enumerate(charges):
                    logger.debug(self, '%.9g    %s', z, coords[i])
            return self

        def get_hcore(self, mol=None):
            if mol is None: mol = self.mol
            if getattr(scf_method, 'get_hcore', None):
                h1e = method_class.get_hcore(self, mol)
            else:  # DO NOT modify post-HF objects to avoid the MM charges applied twice
                raise RuntimeError('mm_charge function cannot be applied on post-HF methods')

            if pyscf.DEBUG:
                v = 0
                for i,q in enumerate(charges):
                    mol.set_rinv_origin(coords[i])
                    v += mol.intor('int1e_rinv') * -q
            else:
                if mol.cart:
                    intor = 'int3c2e_cart'
                else:
                    intor = 'int3c2e_sph'
                nao = mol.nao
                max_memory = self.max_memory - lib.current_memory()[0]
                blksize = int(min(max_memory*1e6/8/nao**2, 200))
                cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas,
                                                     mol._env, intor)
                v = 0
                for i0, i1 in lib.prange(0, charges.size, blksize):
                    fakemol = gto.fakemol_for_charges(coords[i0:i1])
                    j3c = df.incore.aux_e2(mol, fakemol, intor=intor,
                                           aosym='s2ij', cintopt=cintopt)
                    v += numpy.einsum('xk,k->x', j3c, -charges[i0:i1])
                v = lib.unpack_tril(v)
            return h1e + v

        def energy_nuc(self):
# nuclei lattice interaction
            nuc = self.mol.energy_nuc()
            for j in range(self.mol.natm):
                q2, r2 = self.mol.atom_charge(j), self.mol.atom_coord(j)
                r = lib.norm(r2-coords, axis=1)
                nuc += q2*(charges/r).sum()
            return nuc

        def nuc_grad_method(self):
            scf_grad = method_class.nuc_grad_method(self)
            return mm_charge_grad(scf_grad, coords, charges, 'Bohr')

    return QMMM()
add_mm_charges = mm_charge

def mm_charge_grad(scf_grad, coords, charges, unit=None):
    '''Apply the MM charges in the QM gradients' method.  It affects both the
    electronic and nuclear parts of the QM fragment.

    Args:
        scf_grad : a HF or DFT gradient object (grad.HF or grad.RKS etc)
            Once mm_charge_grad function is applied on the SCF object,
            it affects all post-HF calculations eg MP2, CCSD, MCSCF etc
        coords : 2D array, shape (N,3)
            MM particle coordinates
        charges : 1D array
            MM particle charges
    Kwargs:
        unit : str
            Bohr, AU, Ang (case insensitive). Default is the same to mol.unit

    Returns:
        Same gradeints method object as the input scf_grad method

    Examples:

    >>> from pyscf import gto, scf, grad
    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = mm_charge(scf.RHF(mol), [(0.5,0.6,0.8)], [-0.3])
    >>> mf.kernel()
    -101.940495711284
    >>> hfg = mm_charge_grad(grad.hf.RHF(mf), coords, charges)
    >>> hfg.kernel()
    [[-0.25912357 -0.29235976 -0.38245077]
     [-1.70497052 -1.89423883  1.2794798 ]]
    '''
    assert(isinstance(scf_grad, grad.rhf.Gradients))
    if getattr(scf_grad.base, 'with_x2c', None):
        raise NotImplementedError('X2C with QM/MM charges')

    if unit is None:
        unit = scf_grad.mol.unit
    if unit.startswith(('B','b','au','AU')):
        coords = numpy.asarray(coords, order='C')
    elif unit.startswith(('A','a')):
        coords = numpy.asarray(coords, order='C') / lib.parameters.BOHR
    else:
        coords = numpy.asarray(coords, order='C') / unit
    charges = numpy.asarray(charges)

    grad_class = scf_grad.__class__
    class QMMM(grad_class, _QMMMGrad):
        def __init__(self, scf_grad):
            self.__dict__.update(scf_grad.__dict__)

        def dump_flags(self, verbose=None):
            grad_class.dump_flags(self, verbose)
            logger.info(self, '** Add background charges for %s **', grad_class)
            if self.verbose >= logger.DEBUG1:
                logger.debug1(self, 'Charge      Location')
                for i, z in enumerate(charges):
                    logger.debug1(self, '%.9g    %s', z, coords[i])
            return self

        def get_hcore(self, mol=None):
            ''' (QM 1e grad) + <-d/dX i|q_mm/r_mm|j>'''
            if mol is None: mol = self.mol
            g_qm = grad_class.get_hcore(self, mol)
            nao = g_qm.shape[1]
            if pyscf.DEBUG:
                v = 0
                for i,q in enumerate(charges):
                    mol.set_rinv_origin(coords[i])
                    v += mol.intor('int1e_iprinv', comp=3) * q
            else:
                if mol.cart:
                    intor = 'int3c2e_ip1_cart'
                else:
                    intor = 'int3c2e_ip1_sph'
                nao = mol.nao
                max_memory = self.max_memory - lib.current_memory()[0]
                blksize = int(min(max_memory*1e6/8/nao**2, 200))
                cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas,
                                                     mol._env, intor)
                v = 0
                for i0, i1 in lib.prange(0, charges.size, blksize):
                    fakemol = gto.fakemol_for_charges(coords[i0:i1])
                    j3c = df.incore.aux_e2(mol, fakemol, intor, aosym='s1',
                                           comp=3, cintopt=cintopt)
                    v += numpy.einsum('ipqk,k->ipq', j3c, charges[i0:i1])
            return g_qm + v

        def grad_nuc(self, mol=None, atmlst=None):
            if mol is None: mol = self.mol
            g_qm = grad_class.grad_nuc(self, mol, atmlst)
# nuclei lattice interaction
            g_mm = numpy.empty((mol.natm,3))
            for i in range(mol.natm):
                q1 = mol.atom_charge(i)
                r1 = mol.atom_coord(i)
                r = lib.norm(r1-coords, axis=1)
                g_mm[i] = -q1 * numpy.einsum('i,ix,i->x', charges, r1-coords, 1/r**3)
            if atmlst is not None:
                g_mm = g_mm[atmlst]
            return g_qm + g_mm
    return QMMM(scf_grad)

# A tag to label the derived class
class _QMMM:
    pass
class _QMMMGrad:
    pass

# Inject QMMM interface wrapper to other modules
scf.hf.SCF.QMMM = mm_charge
mcscf.casci.CASCI.QMMM = mm_charge
grad.rhf.Gradients.QMMM = mm_charge_grad

if __name__ == '__main__':
    from pyscf import scf, cc, grad
    mol = gto.Mole()
    mol.atom = ''' O                  0.00000000    0.00000000   -0.11081188
                   H                 -0.00000000   -0.84695236    0.59109389
                   H                 -0.00000000    0.89830571    0.52404783 '''
    mol.basis = 'cc-pvdz'
    mol.build()

    coords = [(0.5,0.6,0.8)]
    #coords = [(0.0,0.0,0.0)]
    charges = [-0.5]
    mf = mm_charge(scf.RHF(mol), coords, charges)
    print(mf.kernel()) # -76.3206550372

    g = mf.nuc_grad_method().kernel()
    mfs = mf.as_scanner()
    e1 = mfs(''' O                  0.00100000    0.00000000   -0.11081188
             H                 -0.00000000   -0.84695236    0.59109389
             H                 -0.00000000    0.89830571    0.52404783 ''')
    e2 = mfs(''' O                 -0.00100000    0.00000000   -0.11081188
             H                 -0.00000000   -0.84695236    0.59109389
             H                 -0.00000000    0.89830571    0.52404783 ''')
    print((e1 - e2)/0.002 * lib.param.BOHR, g[0,0])

    mycc = cc.ccsd.CCSD(mf)
    ecc, t1, t2 = mycc.kernel() # ecc = -0.228939687075

    g = mycc.nuc_grad_method().kernel()
    ccs = mycc.as_scanner()
    e1 = ccs(''' O                  0.00100000    0.00000000   -0.11081188
             H                 -0.00000000   -0.84695236    0.59109389
             H                 -0.00000000    0.89830571    0.52404783 ''')
    e2 = ccs(''' O                 -0.00100000    0.00000000   -0.11081188
             H                 -0.00000000   -0.84695236    0.59109389
             H                 -0.00000000    0.89830571    0.52404783 ''')
    print((e1 - e2)/0.002 * lib.param.BOHR, g[0,0])