# -*- coding: utf-8 -*-
"""
Created on Wed Jul 21 11:58:46 2021

@author: wanxiang.shen@u.nus.edu
"""

import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from biopandas.pdb import PandasPdb
from Bio import PDB
import io, PIL
from sklearn.metrics import pairwise_distances

from .agg import AggMolMap

#The hydrophobicity values are from JACS, 1962, 84: 4240-4246. (C. Tanford).
_Hydrophobicity={"A":0.62,"R":-2.53,"N":-0.78,"D":-0.90,"C":0.29,"Q":-0.85,
                 "E":-0.74,"G":0.48,"H":-0.40,"I":1.38,"L":1.06,"K":-1.50,
                 "M":0.64,"F":1.19,"P":0.12,"S":-0.18,"T":-0.05,"W":0.81,"Y":0.26,"V":1.08}

#The hydrophilicity values are from PNAS, 1981, 78:3824-3828 (T.P.Hopp & K.R.Woods).
_hydrophilicity={"A":-0.5,"R":3.0,"N":0.2,"D":3.0,"C":-1.0,"Q":0.2,"E":3.0,
                 "G":0.0,"H":-0.5,"I":-1.8,"L":-1.8,"K":3.0,"M":-1.3,
                 "F":-2.5,"P":0.0,"S":0.3,"T":-0.4,"W":-3.4,"Y":-2.3,"V":-1.5}

#The side-chain mass: CRC Handbook of Chemistry and Physics, 66th ed., CRC Press, Boca Raton, Florida (1985).

_residuemass={"A":15.0,"R":101.0,"N":58.0,"D":59.0,"C":47.0,"Q":72.0,
              "E":73.0,"G":1.000,"H":82.0,"I":57.0,"L":57.0,"K":73.0,
              "M":75.0,"F":91.0,"P":42.0,"S":31.0,"T":45.0,"W":130.0,"Y":107.0,"V":43.0}

#R.M.C. Dawson, D.C. Elliott, W.H. Elliott, K.M. Jones, Data for Biochemical Research 3rd ed., Clarendon Press Oxford (1986).
_pK1={"A":2.35,"C":1.71,"D":1.88,"E":2.19,"F":2.58,"G":2.34,"H":1.78,
      "I":2.32,"K":2.20,"L":2.36,"M":2.28,"N":2.18,"P":1.99,"Q":2.17,
      "R":2.18,"S":2.21,"T":2.15,"V":2.29,"W":2.38,"Y":2.20}

_pK2={"A":9.87,"C":10.78,"D":9.60,"E":9.67,"F":9.24,"G":9.60,"H":8.97,
      "I":9.76,"K":8.90,"L":9.60,"M":9.21,"N":9.09,"P":10.6,"Q":9.13,
      "R":9.09,"S":9.15,"T":9.12,"V":9.74,"W":9.39,"Y":9.11}

_pI={"A":6.11,"C":5.02,"D":2.98,"E":3.08,"F":5.91,"G":6.06,"H":7.64,
     "I":6.04,"K":9.47,"L":6.04,"M":5.74,"N":10.76,"P":6.30,
     "Q":5.65,"R":10.76,"S":5.68,"T":5.60,"V":6.02,"W":5.88,"Y":5.63}

_AvFlexibility={"A":0.357,"R":0.529,"N":0.463,"D":0.511,"C":0.346,"Q":0.493,
                "E":0.497,"G":0.544,"H":0.323,"I":0.462,"L":0.365,"K":0.466,"M":0.295,
                "F":0.314,"P":0.509,"S":0.507,"T":0.444,"W":0.305,"Y":0.420,"V":0.386}

_Polarizability={"A":0.046,"R":0.291,"N":0.134,"D":0.105,"C":0.128,"Q":0.180,
                 "E":0.151,"G":0.000,"H":0.230,"I":0.186,"L":0.186,"K":0.219,
                 "M":0.221,"F":0.290,"P":0.131,"S":0.062,"T":0.108,"W":0.409,"Y":0.298,"V":0.140}

_FreeEnergy={"A":-0.368,"R":-1.03,"N":0.0,"D":2.06,"C":4.53,"Q":0.731,
             "E":1.77,"G":-0.525,"H":0.0,"I":0.791,"L":1.07,"K":0.0,"M":0.656,
             "F":1.06,"P":-2.24,"S":-0.524,"T":0.0,"W":1.60,"Y":4.91,"V":0.401}

_ResidueASA={"A":115.0,"R":225.0,"N":160.0,"D":150.0,"C":135.0,"Q":180.0,
             "E":190.0,"G":75.0,"H":195.0,"I":175.0,"L":170.0,"K":200.0,"M":185.0,
             "F":210.0,"P":145.0,"S":115.0,"T":140.0,"W":255.0,"Y":230.0,"V":155.0}

_ResidueVol={"A":52.6,"R":109.1,"N":75.7,"D":68.4,"C":68.3,"Q":89.7,"E":84.7,
             "G":36.3,"H":91.9,"I":102.0,"L":102.0,"K":105.1,"M":97.7,"F":113.9,
             "P":73.6,"S":54.9,"T":71.2,"W":135.4,"Y":116.2,"V":85.1}

_Steric={"A":0.52,"R":0.68,"N":0.76,"D":0.76,"C":0.62,"Q":0.68,"E":0.68,
         "G":0.00,"H":0.70,"I":1.02,"L":0.98,"K":0.68,"M":0.78,"F":0.70,
         "P":0.36,"S":0.53,"T":0.50,"W":0.70,"Y":0.70,"V":0.76}

_Mutability={"A":100.0,"R":65.0,"N":134.0,"D":106.0,"C":20.0,"Q":93.0,"E":102.0,
             "G":49.0,"H":66.0,"I":96.0,"L":40.0,"K":-56.0,"M":94.0,"F":41.0,"P":56.0,
             "S":120.0,"T":97.0,"W":18.0,"Y":41.0,"V":74.0}

def standard_scale(aap):
    scaler = StandardScaler()
    s = pd.Series(aap)
    res = scaler.fit_transform(s.values.reshape(-1,1)).reshape(-1,)
    return pd.Series(res, index=s.index).to_dict()

IntrinsicAAPs = {'Hydrophobicity':_Hydrophobicity,
                'Hydrophilicity':_hydrophilicity,
                'ASA':_ResidueASA,
                'Flexibility': _AvFlexibility,
                'FreeEnergy': _FreeEnergy,
                'Steric': _Steric,
                'pKa':_pK1,
                'pKb':_pK2,
                'pI':_pI}

class IntrinsicAAP:
    def __init__(self, **entries):
        self.__dict__.update(entries)

def get_pdb_xyzb_ca(df_aa):

    '''
    https://en.wikipedia.org/wiki/Protein_contact_map:
    df_aa is the dataframe that is groupbyed from `residue_name` and  `residue_number`
    '''
    #df_aa = df[(df['residue_name'] == 'PRO') & (df['residue_number'] ==2)]

    ts = df_aa[df_aa.atom_name == 'CA']
    if len(ts == 1):
        ts = ts.iloc[-1] ## if multi-CA, select the last one
        x,y,z,b = ts.x_coord, ts.y_coord, ts.z_coord, ts.b_factor
    else:
        x,y,z,b = np.nan, np.nan, np.nan, np.nan
        aa = df_aa[['residue_name', 'residue_number']].iloc[0].to_dict()
        print('CA atom not exists in the residue: %s' % aa)
    return x,y,z,b

def get_pdb_xyzb_cb(df_aa):
    '''
    https://en.wikipedia.org/wiki/Protein_contact_map:

    df_aa is the dataframe that is groupbyed from `residue_name` and  `residue_number`
    '''
    #df_aa = df[(df['residue_name'] == 'PRO') & (df['residue_number'] ==2)]

    ts = df_aa[df_aa.atom_name == 'CB'] ## if multi-CA, select the last one

    if len(ts) ==1:
        ts = ts.iloc[-1]
        x,y,z,b = ts.x_coord, ts.y_coord, ts.z_coord, ts.b_factor
    else:
        ##distance between Cβ-Cβ atoms with threshold 6-12 Å (Cα is used for Glycine):
        ts = df_aa[df_aa.atom_name == 'CA'] #Case Glycine
        if len(ts) ==1:
            ts = ts.iloc[-1]
            x,y,z,b = ts.x_coord, ts.y_coord, ts.z_coord, ts.b_factor
        else:
            x,y,z,b = np.nan, np.nan, np.nan, np.nan
            aa = df_aa[['residue_name', 'residue_number']].iloc[0].to_dict()
            print('CA atom not exists in the residue: %s' % aa)
    return x,y,z,b

def get_pdb_xyzb_mean(df_aa):

    '''
    get the mean coord.
    https://en.wikipedia.org/wiki/Protein_contact_map:
    df_aa is the dataframe that is groupbyed from `residue_name` and  `residue_number`
    '''
    #df_aa = df[(df['residue_name'] == 'PRO') & (df['residue_number'] ==2)]

    x,y,z,b = df_aa[['x_coord', 'y_coord', 'z_coord', 'b_factor']].mean().tolist()

    return x,y,z,b

class PDB2Fmap:

    def __init__(self, embd_grain = 'CA', fmap_shape = None):

        '''
        embd_grain: {'CA', 'CB', 'mean', 'all'}
        '''
        self.embd_grain = embd_grain
        self.fmap_shape = fmap_shape

    def fit(self, pdb_file, embd_chain = None):
        '''
        pdb_file: pdf file path
        embd_chain: pdb chain to do embedding
        '''
        self.pdb_file = pdb_file
        self.pdb = PandasPdb().read_pdb(self.pdb_file)
        self.embd_chain = embd_chain

        if embd_chain != None:
            self.dfpdb = self.pdb.df['ATOM'][self.pdb.df['ATOM'].chain_id == embd_chain]
        else:
            self.dfpdb = self.pdb.df['ATOM']
        if self.embd_grain == 'mean':
            df_embd = self.dfpdb.groupby(['residue_number', 'residue_name']).apply(get_pdb_xyzb_mean).apply(pd.Series)
            df_embd.columns = ['x_coord', 'y_coord', 'z_coord', 'b_factor']
            df_embd = df_embd.reset_index()

        if self.embd_grain == 'CB':
            df_embd = self.dfpdb.groupby(['residue_number', 'residue_name']).apply(get_pdb_xyzb_cb).apply(pd.Series)
            df_embd.columns = ['x_coord', 'y_coord', 'z_coord', 'b_factor']
            df_embd = df_embd.reset_index()

        if self.embd_grain == 'CA':
            df_embd = self.dfpdb.groupby(['residue_number', 'residue_name']).apply(get_pdb_xyzb_ca).apply(pd.Series)
            df_embd.columns = ['x_coord', 'y_coord', 'z_coord', 'b_factor']
            df_embd = df_embd.reset_index()

        if self.embd_grain == 'all':
            df_embd = self.dfpdb[['residue_name', 'residue_number', 'x_coord', 'y_coord','z_coord', 'b_factor']]

        df_embd['residue_name_1aa'] = df_embd['residue_name'].map(PDB.protein_letters_3to1)
        df_embd.index = df_embd.index.astype(str) + '-' + df_embd['residue_name_1aa']
        dfx = df_embd[['x_coord','y_coord','z_coord']].T
        self.dfx = dfx
        self.df_embd = df_embd
        self.mp = AggMolMap(dfx, metric='euclidean')
        self.mp.fit(fmap_shape = self.fmap_shape, cluster_channels=1)
        self.fmap_shape = self.mp.fmap_shape

    def transform_xyz(self, scale = True, feature_range=(0, 1)):
        '''
        x, y, z coordinates
        '''
        if scale:
            scaler = MinMaxScaler(feature_range = feature_range)
            x = scaler.fit_transform(self.dfx.T).T
        else:
            x = self.dfx.values
        X = self.mp.batch_transform(x, scale=False)
        return X

    def transofrm_bf(self, scale = True, feature_range=(0, 1)):
        '''
        b-factor
        '''
        if scale:
            scaler = MinMaxScaler(feature_range = feature_range)
            x = scaler.fit_transform(self.df_embd[['b_factor']]).T
        else:
            x = self.df_embd[['b_factor']].values.T
        X = self.mp.transform(x[0], scale=False)
        return X

    def transofrm_pkt(self, pkt_file):
        '''
        pocket pdb file
        '''
        self.pkt_file = pkt_file

        ## pocket
        self.pkt = PandasPdb().read_pdb(self.pkt_file)
        if self.embd_chain != None:
            self.dfpkt = self.pkt.df['ATOM'][self.pkt.df['ATOM'].chain_id == self.embd_chain]
        else:
            self.dfpkt = self.pkt.df['ATOM']

        pkt_residue_number = self.dfpkt.residue_number.unique()

        self.df_embd['pocket'] = self.df_embd.residue_number.isin(pkt_residue_number)*1
        x = self.df_embd[['pocket']].values.T
        X = self.mp.transform(x[0], scale=False)
        return X

    def transform_custom(self, aap_df, scale = True, feature_range=(0, 1)):

        '''
        aap_df: dataframe of animo acid propetries, each column is one type of property, total 20 rows for all 20 types of animo acids
        aap_df example:
        ==============
        >>> from molmap.feature.sequence.aas.local_feature.aai import load_index
        >>> aap_df = load_index.data.T
        '''

        df_custom = pd.DataFrame(index = self.df_embd.index)
        for k, v in aap_df.to_dict().items():
            df_custom[k] = self.df_embd.residue_name_1aa.map(v)
        self.df_custom = df_custom
        if scale:
            scaler = MinMaxScaler(feature_range = feature_range)
            x = scaler.fit_transform(self.df_custom).T
        else:
            x = self.df_custom.values.T
        X = self.mp.batch_transform(x, scale=False)

        return X

    def transform_intrinsic(self, scale = True, feature_range=(0, 1)):

        df_intrinsic = pd.DataFrame(index = self.df_embd.index)
        for k, v in IntrinsicAAPs.items():
            df_intrinsic[k] = self.df_embd.residue_name_1aa.map(v)
        self.df_intrinsic = df_intrinsic
        if scale:
            scaler = MinMaxScaler(feature_range = feature_range)
            x = scaler.fit_transform(self.df_intrinsic).T
        else:
            x = self.df_intrinsic.values.T
        X = self.mp.batch_transform(x, scale=False)

        return X

class PDB2Img:

    def __init__(self, pdb_file,  embd_grain = 'CA', embd_chain = None):

        '''
        embd_grain: {'CA', 'CB', 'mean', 'all'}
        pdb_file: pdf file path
        embd_chain: pdb chain to do embedding
        '''
        self.embd_grain = embd_grain
        self.pdb_file = pdb_file
        self.pdb = PandasPdb().read_pdb(self.pdb_file)
        self.embd_chain = embd_chain

        if embd_chain != None:
            self.dfpdb = self.pdb.df['ATOM'][self.pdb.df['ATOM'].chain_id == embd_chain]
        else:
            self.dfpdb = self.pdb.df['ATOM']
        if self.embd_grain == 'mean':
            df_embd = self.dfpdb.groupby(['residue_number', 'residue_name']).apply(get_pdb_xyzb_mean).apply(pd.Series)
            df_embd.columns = ['x_coord', 'y_coord', 'z_coord', 'b_factor']
            df_embd = df_embd.reset_index()

        if self.embd_grain == 'CB':
            df_embd = self.dfpdb.groupby(['residue_number', 'residue_name']).apply(get_pdb_xyzb_cb).apply(pd.Series)
            df_embd.columns = ['x_coord', 'y_coord', 'z_coord', 'b_factor']
            df_embd = df_embd.reset_index()

        if self.embd_grain == 'CA':
            df_embd = self.dfpdb.groupby(['residue_number', 'residue_name']).apply(get_pdb_xyzb_ca).apply(pd.Series)
            df_embd.columns = ['x_coord', 'y_coord', 'z_coord', 'b_factor']
            df_embd = df_embd.reset_index()

        if self.embd_grain == 'all':
            df_embd = self.dfpdb[['residue_name', 'residue_number', 'x_coord', 'y_coord','z_coord', 'b_factor']]

        df_embd['residue_name_1aa'] = df_embd['residue_name'].map(PDB.protein_letters_3to1)
        df_embd.index = df_embd.index.astype(str) + '-' + df_embd['residue_name_1aa']
        dfx = df_embd[['x_coord','y_coord','z_coord']]
        self.dfx = dfx
        self.df_embd = df_embd
        self.dfx_dist = pairwise_distances(self.dfx)

    def transform(self, fmap_shape = None, cmap='jet_r', vmin = 0, vmax=80, dpi=100):
        '''
        fig size: dpi*3
        '''
        fig, ax = plt.subplots()
        ax.imshow(self.dfx_dist, cmap=cmap, vmin = vmin, vmax=vmax)
        ax.axis('off')
        with io.BytesIO() as buff:
            fig.savefig(buff, bbox_inches='tight', pad_inches=0, dpi=dpi)
            buff.seek(0)
            im = PIL.Image.open(buff)
            im = im.convert('RGB')
            if fmap_shape != None:
                im = im.resize(fmap_shape)
            x = np.array(im) / 255

        return fig, x

if __name__ == '__main__':
    pm = PDB2Fmap(embd_grain='all', fmap_shape=None)
    pm.fit(pdb_file='./1a1e/1a1e_protein.pdb', embd_chain='B')
    X = pm.transform_xyz(scale=True, feature_range=(0.1,1))
    X = pm.transofrm_bf(scale = True, feature_range=(0.2,1))
    X = pm.transofrm_pkt('./1a1e/1a1e_pocket.pdb')
    X = pm.transform_intrinsic()
    sns.heatmap(X[2].reshape(*pm.fmap_shape), cmap = 'jet')


r'''Custom Printers for Sympy expressions

These classes are used by default by the QAlgebra printing systems as
sub-printers for SymPy objects (e.g. for symbolic coefficients). They fix some
issues with SymPy's builtin printers:

* factors like $\frac{1}{\sqrt{2}}$ occur very commonly in quantum mechanics,
  and it is standard notation to write them as such. SymPy insists on
  rationalizing denominators, using $\frac{\sqrt{2}}{2}$ instead. Our custom
  printers restore the canonical form. Note that internally, Sympy still uses
  the rationalized structure; but in any case, Sympy makes no guarantees
  between the algebraic structure of an expression and how it is printed.
* Symbols (especially greek letters) are extremely common, and it's much more
  readable if the string representation of an expression uses unicode for
  these. SymPy supports unicode "pretty-printing"
  (:func:`sympy.printing.pretty.pretty.pretty_print`) only in "2D", where
  expressions are rendered as multiline unicode strings. While this is fine for
  interactive display, it does not work so well for a simple ``str``. The
  :class:`SympyUnicodePrinter` solves this by producing simple strings with
  unicode symbols.
* Some algebraic structures such as factorials, complex-conjugates and indexed
  symbols have sub-optimal rendering in :class:`sympy.printing.str.StrPrinter`
* QAlgebra contains some custom subclasses of SymPy objects (e.g.
  :class:`.IdxSym`) that the default printers don't know how to deal with
  (respectively, render incorrectly!)
'''
from collections import OrderedDict

import sympy
from sympy import sqrt
from sympy.core import Mul, Pow, Rational, S
from sympy.core.mul import _keep_coeff
from sympy.printing.latex import LatexPrinter
from sympy.printing.precedence import PRECEDENCE, precedence
from sympy.printing.pretty.pretty_symbology import pretty_symbol
from sympy.printing.repr import ReprPrinter
from sympy.printing.str import StrPrinter

from ._unicode_mappings import _SUBSCRIPT_MAPPING, _SUPERSCRIPT_MAPPING

__all__ = []
__private__ = [
    'SympyLatexPrinter',
    'SympyStrPrinter',
    'SympyUnicodePrinter',
    'SympyReprPrinter',
    'derationalize_denom',
]

delattr(sympy.Indexed, '_sympystr')  # not sure how else to override printer

def derationalize_denom(expr):
    """Try to de-rationalize the denominator of the given expression.

    The purpose is to allow to reconstruct e.g. ``1/sqrt(2)`` from
    ``sqrt(2)/2``.

    Specifically, this matches `expr` against the following pattern::

        Mul(..., Rational(n, d), Pow(d, Rational(1, 2)), ...)

    and returns a tuple ``(numerator, denom_sq, post_factor)``, where
    ``numerator`` and ``denom_sq`` are ``n`` and ``d`` in the above pattern (of
    type `int`), respectively, and ``post_factor`` is the product of the
    remaining factors (``...`` in `expr`). The result will fulfill the
    following identity::

        (numerator / sqrt(denom_sq)) * post_factor == expr

    If `expr` does not follow the appropriate pattern, a :exc:`ValueError` is
    raised.
    """
    r_pos = -1
    p_pos = -1
    numerator = S.Zero
    denom_sq = S.One
    post_factors = []
    if isinstance(expr, Mul):
        for pos, factor in enumerate(expr.args):
            if isinstance(factor, Rational) and r_pos < 0:
                r_pos = pos
                numerator, denom_sq = factor.p, factor.q
            elif isinstance(factor, Pow) and r_pos >= 0:
                if factor == sqrt(denom_sq):
                    p_pos = pos
                else:
                    post_factors.append(factor)
            else:
                post_factors.append(factor)
        if r_pos >= 0 and p_pos >= 0:
            return numerator, denom_sq, Mul(*post_factors)
        else:
            raise ValueError("Cannot derationalize")
    else:
        raise ValueError("expr is not a Mul instance")

class SympyStrPrinter(StrPrinter):
    """Variation of sympy ``StrPrinter`` that derationalizes denominators.

    Additionally, it contains the following modifications:

    * Support for :class:`.IdxSym`
    * Rendering of :class:`sympy.tensor.indexed.Indexed` as subscripts
    * Rendering of :class:`sympy.functions.combinatorial.factorials.factorial`
      as ``!``
    * Option `conjg_style` to configure how complex conjugates are rendered:
      ``'func' renders it as ``conjugate(...)``, and ``'star'`` uses an
      exponentiated asterisk
    """

    printmethod = "_sympystr"
    _default_settings = {
        "order": None,
        "full_prec": "auto",
        "sympy_integers": False,
        "conjg_style": 'func',
    }

    # _print_IdxSym(self, expr) is implemented in IdxSym._sympystr

    def _print_Mul(self, expr):

        prec = precedence(expr)

        try:
            numerator, denom_sq, post_factor = derationalize_denom(expr)
            if post_factor == S.One:
                return "%s/sqrt(%s)" % (numerator, denom_sq)
            else:
                if numerator == 1:
                    return "%s / sqrt(%s)" % (
                        self.parenthesize(post_factor, prec),
                        denom_sq,
                    )
                else:
                    return "(%s/sqrt(%s)) %s" % (
                        numerator,
                        denom_sq,
                        self.parenthesize(post_factor, prec),
                    )
        except ValueError:
            return super()._print_Mul(expr)

    def _print_Indexed(self, expr):
        indices = [self._print(i) for i in expr.indices]
        sep = ','
        if all([len(i.replace("'", '')) == 1 for i in indices]):
            sep = ''
        return self._print(expr.base) + '_%s' % sep.join(indices)

    def _print_factorial(self, expr, exp=None):
        res = r"%s!" % self.parenthesize(expr.args[0], PRECEDENCE["Func"])

        if exp is not None:
            return r"%s^{%s}" % (res, exp)
        else:
            return res

    def _print_conjugate(self, expr, exp=None):
        if self._settings['conjg_style'] == 'star':
            res = self.parenthesize(expr.args[0], PRECEDENCE["Func"]) + '^*'
        elif self._settings['conjg_style'] in ['func', 'overbar']:
            # recognizing "overbar" is just for compatibility with the other
            # printers
            res = r'conjugate(' + self._print(expr.args[0]) + r')'
            pass
        else:
            raise ValueError(
                "The 'conjg_style' setting must be one of " "'star', 'func'"
            )
        if exp is not None:
            return r"%s^%s" % (res, exp)
        else:
            return res

class SympyLatexPrinter(LatexPrinter):
    """Variation of sympy ``LatexPrinter`` that derationalizes denominators.

    Additionally, it contains the following modifications:

    * Support for :class:`.IdxSym`
    * A setting `conjg_style` that allows to specify how complex conjugate are
      rendered: ``'overline'`` (the default) draws a line over the number,
      'star' uses an exponentiated asterisk, and 'func' renders a a
      ``conjugate`` function
    """

    printmethod = "_latex"

    _default_settings = LatexPrinter._default_settings.copy()
    _default_settings.update(
        {
            "order": None,
            "mode": "plain",
            "itex": False,
            "fold_frac_powers": False,
            "fold_func_brackets": False,
            "fold_short_frac": None,
            "long_frac_ratio": 2,
            "mul_symbol": None,
            "inv_trig_style": "abbreviated",
            "mat_str": None,
            "mat_delim": "[",
            "symbol_names": {},
            "conjg_style": 'overline',
        }
    )

    # _print_IdxSym(self, expr) is implemented in IdxSym._latex

    def _print_Mul(self, expr):

        prec = precedence(expr)

        try:
            numerator, denom_sq, post_factor = derationalize_denom(expr)
            if post_factor == S.One:
                return r'\frac{%s}{\sqrt{%s}}' % (numerator, denom_sq)
            else:
                if numerator == 1:
                    return r'\frac{%s}{\sqrt{%s}}' % (
                        self._print(post_factor),
                        denom_sq,
                    )
                else:
                    return r'\frac{%s}{\sqrt{%s}} %s' % (
                        numerator,
                        denom_sq,
                        self.parenthesize(post_factor, prec),
                    )
        except ValueError:
            return super()._print_Mul(expr)

    def _print_conjugate(self, expr, exp=None):
        if self._settings['conjg_style'] == 'overline':
            tex = r"\overline{%s}" % self._print(expr.args[0])
        elif self._settings['conjg_style'] == 'star':
            tex = r"{%s}^*" % self.parenthesize(
                expr.args[0], PRECEDENCE["Func"]
            )
        elif self._settings['conjg_style'] == 'func':
            tex = (
                r'\operatorname{conjugate}\left('
                + self._print(expr.args[0])
                + r'\right)'
            )
        else:
            raise ValueError(
                "The 'conjg_style' setting must be one of "
                "'overline', 'star', 'func'"
            )

        if exp is not None:
            return r"{%s}^{%s}" % (tex, exp)
        else:
            return tex

    def _print_Indexed(self, expr):
        from qalgebra.printing.latexprinter import _TEX_SINGLE_LETTER_SYMBOLS

        indices = [self._print(i) for i in expr.indices]
        sep = ','
        all_indices_are_single_letters = True
        for i in indices:
            i = i.replace(r'\prime', '')
            if (i not in _TEX_SINGLE_LETTER_SYMBOLS) and (len(i) > 1):
                all_indices_are_single_letters = False
        if all_indices_are_single_letters:
            sep = ' '
        return self._print(expr.base) + '_{%s}' % sep.join(indices)

class SympyUnicodePrinter(SympyStrPrinter):
    """Printer that represents SymPy expressions as (single-line) unicode
    strings.

    This is a mixture of SymPy's ``StrPrinter`` and
    :class:`sympy.printing.pretty.pretty.PrettyPrinter` (minus the 2D
    printing), with the same extensions as :class:`SympyStrPrinter`
    """

    printmethod = "_sympystr"
    _default_settings = {
        "order": None,
        "full_prec": "auto",
        "sympy_integers": False,
        "superscript_asterisk_sym": "\u00A0\u20F0",
        "conjg_style": 'star',
    }

    def _print_Add(self, expr, order=None):
        if self.order == 'none':
            terms = list(expr.args)
        else:
            terms = self._as_ordered_terms(expr, order=order)

        PREC = precedence(expr)
        l = []
        for term in terms:
            t = str(self._print(term))
            if t.startswith('-'):
                sign = "-"
                t = t[1:]
            else:
                sign = "+"
            if precedence(term) < PREC:
                l.extend([sign, "(%s)" % t])
            else:
                l.extend([sign, t])
        sign = l.pop(0)
        if sign == '+':
            sign = ""
        return sign + ' '.join(l)

    def _print_ComplexInfinity(self, expr):
        return '∞'

    def _print_ImaginaryUnit(self, expr):
        return 'ⅈ'

    def _print_Infinity(self, expr):
        return '∞'

    def _print_Inverse(self, I):
        return "%s⁻¹" % self.parenthesize(I.arg, PRECEDENCE["Pow"])

    def _print_Pi(self, expr):
        return 'π'

    def _print_Mul(self, expr):

        prec = precedence(expr)

        try:
            numerator, denom_sq, post_factor = derationalize_denom(expr)
            if post_factor == S.One:
                return "%s/√%s" % (numerator, denom_sq)
            else:
                if numerator == 1:
                    return "%s / √%s" % (
                        self.parenthesize(post_factor, prec),
                        denom_sq,
                    )
                else:
                    return "(%s/√%s) %s" % (
                        numerator,
                        denom_sq,
                        self.parenthesize(post_factor, prec),
                    )
        except ValueError:
            pass  # Continue below

        c, e = expr.as_coeff_Mul()
        if c < 0:
            expr = _keep_coeff(-c, e)
            sign = "-"
        else:
            sign = ""

        a = []  # items in the numerator
        b = []  # items that are in the denominator (if any)

        if self.order not in ('old', 'none'):
            args = expr.as_ordered_factors()
        else:
            # use make_args in case expr was something like -x -> x
            args = Mul.make_args(expr)

        # Gather args for numerator/denominator
        for item in args:
            if (
                item.is_commutative
                and item.is_Pow
                and item.exp.is_Rational
                and item.exp.is_negative
            ):
                if item.exp != -1:
                    b.append(Pow(item.base, -item.exp, evaluate=False))
                else:
                    b.append(Pow(item.base, -item.exp))
            elif item.is_Rational and item is not S.Infinity:
                if item.p != 1:
                    a.append(Rational(item.p))
                if item.q != 1:
                    b.append(Rational(item.q))
            else:
                a.append(item)

        a = a or [S.One]

        a_str = [str(self.parenthesize(x, prec)) for x in a]
        b_str = [str(self.parenthesize(x, prec)) for x in b]

        if len(b) == 0:
            return sign + ' '.join(a_str)
        elif len(b) == 1:
            return sign + ' '.join(a_str) + "/" + b_str[0]
        else:
            return sign + ' '.join(a_str) + "/(%s)" % ' '.join(b_str)

    def _print_NegativeInfinity(self, expr):
        return '-∞'

    def _print_Pow(self, expr, rational=False):
        PREC = precedence(expr)

        if expr.exp is S.Half and not rational:
            return "√%s" % self.parenthesize(expr.base, PREC)

        if expr.is_commutative:
            if -expr.exp is S.Half and not rational:
                # Note: Don't test "expr.exp == -S.Half" here, because that
                # will match -0.5, which we don't want.
                return "1/√%s" % self.parenthesize(expr.base, PREC)
            if expr.exp is -S.One:
                # Similarly to the S.Half case, don't test with "==" here.
                return '1/%s' % self.parenthesize(expr.base, PREC)

        e = self.parenthesize(expr.exp, PREC)
        if (
            self.printmethod == '_sympyrepr'
            and expr.exp.is_Rational
            and expr.exp.q != 1
        ):
            # the parenthesized exp should be '(Rational(a, b))' so strip
            # parens, but just check to be sure.
            if e.startswith('(Rational'):
                return '%s**%s' % (self.parenthesize(expr.base, PREC), e[1:-1])
        try:
            e_super = ''.join([_SUPERSCRIPT_MAPPING[l] for l in e])
            return '%s%s' % (self.parenthesize(expr.base, PREC), e_super)
        except KeyError:
            return '%s**%s' % (self.parenthesize(expr.base, PREC), e)

    def _print_MatPow(self, expr):
        PREC = precedence(expr)
        b = str(self.parenthesize(expr.base, PREC))
        e = str(self.parenthesize(expr.exp, PREC))
        try:
            e_super = ''.join([_SUPERSCRIPT_MAPPING[l] for l in e])
            return '%s%s' % (b, e_super)
        except KeyError:
            return '%s**%s' % (b, e)

    def _print_Symbol(self, e):
        return pretty_symbol(e.name)

    _print_RandomSymbol = _print_Symbol

    def _print_Identity(self, expr):
        return "𝟙"

    def _print_ZeroMatrix(self, expr):
        return "𝟘"

    def _print_Indexed(self, expr):
        indices = [self._print(i) for i in expr.indices]
        sep = ','
        if all([len(i.replace("'", '')) == 1 for i in indices]):
            sep = ''
        subscript = sep.join(indices)
        try:
            subscript = ''.join([_SUBSCRIPT_MAPPING[l] for l in subscript])
            return self._print(expr.base) + subscript
        except KeyError:
            return self._print(expr.base) + '_%s' % subscript

    def _print_conjugate(self, expr, exp=None):
        if self._settings['conjg_style'] == 'star':
            rendered_arg = self.parenthesize(expr.args[0], PRECEDENCE["Func"])
            if '_' in rendered_arg or '^' in rendered_arg:
                if not rendered_arg.endswith(')'):
                    rendered_arg = "(%s)" % rendered_arg
            res = rendered_arg + self._settings['superscript_asterisk_sym']
        elif self._settings['conjg_style'] in ['func', 'overbar']:
            # recognizing "overbar" is just for compatibility with the other
            # printers
            res = r'conjugate(' + self._print(expr.args[0]) + r')'
            pass
        else:
            raise ValueError(
                "The 'conjg_style' setting must be one of " "'star', 'func'"
            )
        if exp is not None:
            return r"%s^%s" % (res, exp)
        else:
            return res

class SympyReprPrinter(ReprPrinter):
    """Representation printer with support for
    :class:`.IdxSym`"""

    # _print_IdxSym(self, expr) is implemented in IdxSym._sympyrepr

    def _print_Symbol(self, expr):
        d = expr._assumptions.generator

        d = OrderedDict([(key, d[key]) for key in sorted(d.keys())])
        # the use of an OrderedDict is the only diffference between this and
        # ReprPrinter._print_Symbol. It ensures that we always get the same
        # repr

        # print the dummy_index like it was an assumption
        if expr.is_Dummy:
            d['dummy_index'] = expr.dummy_index

        if d == {}:
            return "%s(%s)" % (expr.__class__.__name__, self._print(expr.name))
        else:
            attr = ['%s=%s' % (k, v) for k, v in d.items()]
            return "%s(%s, %s)" % (
                expr.__class__.__name__,
                self._print(expr.name),
                ', '.join(attr),
            )

        res = self._print_Symbol(expr)
        if expr.primed > 0:
            res = res[:-1] + ", primed=%d)" % expr.primed
        return res


#! /usr/bin/python3
# -*- coding: utf-8 -*-

"""
    ****************
    Celestial Target
    ****************

    .. inheritance-diagram:: nenupy.astro.target.FixedTarget nenupy.astro.target.SolarSystemTarget
        :parts: 3

    .. autosummary::

        ~FixedTarget
        ~SolarSystemTarget

"""

from __future__ import annotations

__author__ = "Alan Loh"
__copyright__ = "Copyright 2021, nenupy"
__credits__ = ["Alan Loh"]
__maintainer__ = "Alan"
__email__ = "alan.loh@obspm.fr"
__status__ = "Production"
__all__ = [
    "Target",
    "FixedTarget",
    "SolarSystemTarget"
]

from abc import ABC, abstractmethod
from typing import Callable
import numpy as np
import logging
log = logging.getLogger(__name__)

import astropy.units as u
from astropy.time import Time, TimeDelta
from astropy.coordinates import SkyCoord, EarthLocation, FK5, AltAz

from nenupy import nenufar_position
from nenupy.astro import common_sources
from nenupy.astro.astro_tools import (
    AstroObject,
    hour_angle,
    solar_system_source
)

# ============================================================= #
# -------------------------- Target --------------------------- #
# ============================================================= #
class Target(AstroObject, ABC):
    """ Abstract class to handle target objects.

        .. versionadded:: 2.0.0

        .. rubric:: Attributes Summary

        .. autosummary::

            ~Target.coordinates
            ~Target.time
            ~Target.observer
            ~Target.is_circumpolar
            ~Target.culmination_azimuth

        .. rubric:: Methods Summary

        .. autosummary::

            ~Target.meridian_transit
            ~Target.next_meridian_transit
            ~Target.previous_meridian_transit
            ~Target.azimuth_transit
            ~Target.rise_time
            ~Target.next_rise_time
            ~Target.previous_rise_time
            ~Target.set_time
            ~Target.next_set_time
            ~Target.previous_set_time

        .. rubric:: Attributes and Methods Documentation

    """

    def __init__(self,
            coordinates: SkyCoord,
            observer: EarthLocation = nenufar_position,
            time: Time = Time.now()
        ):
        self.coordinates = coordinates
        self.observer = observer
        self.time = time

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def time(self) -> Time:
        """ """
        return self._time
    @time.setter
    def time(self, t):
        if t.isscalar:
            t = t.reshape(1,)
        self._time = t

    @property
    def is_circumpolar(self) -> bool:
        r""" Whether the celestial object is circumpolar at the
            observer's latitude.

            .. math::
                l + \delta \geq 90\,{\rm deg}

            where :math:`l` is the latitude (defined in
            :attr:`~nenupy.astro.astro_tools.AstroObject.observer`),
            :math:`\delta` is the object's declination (defined in
            :attr:`~nenupy.astro.astro_tools.AstroObject.coordinates`).
        """
        return np.all((self.observer.lat + self.coordinates.dec) >= 90*u.deg)

    @property
    def culmination_azimuth(self) -> u.Quantity:
        """ """
        if not self.is_circumpolar:
            return 180*u.deg
        else:
            return 0*u.deg

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    def meridian_transit(self,
            t_min: Time = Time.now(),
            duration: TimeDelta = TimeDelta(86400, format='sec'),
            precision: TimeDelta = TimeDelta(5, format='sec'),
            fast_compute: bool = True
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` meridian transit time(s).
            This method returns all the transit times found in the time
            window ranging from ``t_min`` to ``t_min + duration``.

            :param t_min:
                Starting time of the temporal window within which
                meridian transits are looked for.
                Default is current time.
            :type t_min:
                :class:`~astropy.time.Time`
            :param duration:
                Width of the temporal window within which
                meridian transits are looked for.
                Default is ``1 day``.
            :type duration:
                :class:`~astropy.time.TimeDelta`
            :param precision:
                Temporal precision of the returned meridian transit values.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`
            :param fast_compute:
                If set to ``True``, a fast approximation is used during
                the computation of Local Sidereal Time.
                Default is ``True``.
            :type fast_compute:
                `bool`

            :returns:
                Meridian transit times.
                If no transit times are found (because the requested
                time window doesn't contain any) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.meridian_transit(
                        t_min=Time("2021-01-01"),
                        duration=TimeDelta(86400*2, format="sec")
                    )
                <Time object: scale='utc' format='iso' value=['2021-01-01 13:05:47.868' '2021-01-02 13:01:51.882']>

            .. seealso::
                :ref:`ephemerides_sec`

        """
        def find_ha_transit(times: Time):
            """ """
            fk5 = self._get_source_coordinates(
                time=times
            ).transform_to(FK5(equinox=times))
            ha = hour_angle(
                radec=fk5,
                time=times,
                observer=self.observer,
                fast_compute=fast_compute
            )
            return np.where(
                (np.roll(ha, shift=-1, axis=1) - ha)[:, :-1] < 0
            )
        return self._find_crossing_times(
            finding_function=find_ha_transit,
            t_min=t_min,
            duration=duration,
            precision=precision
        )

    def azimuth_transit(self,
            azimuth: u.Quantity = 180*u.deg,
            t_min: Time = Time.now(),
            duration: TimeDelta = TimeDelta(86400, format='sec'),
            precision: TimeDelta = TimeDelta(5, format='sec'),
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` transit time(s) at a given ``azimuth`` value.
            This method returns all the transit times found in the time
            window ranging from ``t_min`` to ``t_min + duration``.

            :param azimuth:
                Azimuth at which the transit is computed.
                Default is ``180 deg`` (i.e. South).
            :type azimuth:
                :class:`~astropy.units.Quantity`
            :param t_min:
                Starting time of the temporal window within which
                azimuth transits are looked for.
                Default is current time.
            :type t_min:
                :class:`~astropy.time.Time`
            :param duration:
                Width of the temporal window within which
                azimuth transits are looked for.
                Default is ``1 day``.
            :type duration:
                :class:`~astropy.time.TimeDelta`
            :param precision:
                Temporal precision of the returned azimuth transit values.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`

            :returns:
                Azimuth transit times.
                If no transit times are found (either because the requested
                time window doesn't contain any or because the source apparent
                sky position does not cross the desired ``azimuth``) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> import astropy.units as u
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.azimuth_transit(
                        azimuth=100*u.deg,
                        t_min=Time("2021-01-01"),
                        duration=TimeDelta(86400*2, format="sec")
                    )
                <Time object: scale='utc' format='iso' value=['2021-01-01 11:22:12.463' '2021-01-02 11:18:16.477']>

            .. seealso::
                :ref:`ephemerides_sec`

        """
        def find_az_transit(times: Time):
            """ """
            # altaz_coordinates = radec_to_altaz(
            #     radec=self._get_source_coordinates(time=times),
            #     time=times,
            #     observer=self.observer,
            #     fast_compute=fast_compute
            # ).reshape(times.shape)
            altaz_coordinates = self._get_source_coordinates(time=times).transform_to(
                AltAz(
                    obstime=times,
                    location=self.observer
                )
            )
            azimuths = altaz_coordinates.az.rad
            az = azimuth.to(u.rad).value
            if self.is_circumpolar:
                complexAzStarts = np.angle(
                    np.cos(azimuths[:, :-1]) + 1j*np.sin(azimuths[:, :-1])
                )
                complexAzStops = np.angle(
                    np.cos(azimuths[:, 1:]) + 1j*np.sin(azimuths[:, 1:])
                )
                mask = (complexAzStarts <= az) &\
                    (complexAzStops >= az)
                mask |= (complexAzStarts >= az) &\
                    (complexAzStops <= az)
            else:
                mask = (azimuths[:, :-1] <= az) &\
                    (azimuths[:, 1:] >= az)
            return np.where(mask)
        return self._find_crossing_times(
            finding_function=find_az_transit,
            t_min=t_min,
            duration=duration,
            precision=precision
        )

    def next_meridian_transit(self,
            time: Time = Time.now(),
            precision: TimeDelta = TimeDelta(5, format='sec'),
            fast_compute: bool = True
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` next meridian transit time.
            This method returns the next transit time found after ``time``.

            :param time:
                Relative time used to searching for the next meridian transit.
            :type time:
                :class:`~astropy.time.Time`
            :param precision:
                Temporal precision of the returned meridian transit value.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`
            :param fast_compute:
                If set to ``True``, a fast approximation is used during
                the computation of Local Sidereal Time.
                Default is ``True``.
            :type fast_compute:
                `bool`

            :returns:
                Next meridian transit time.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.next_meridian_transit(
                        time=Time("2021-01-01 12:00:00")
                    )
                <Time object: scale='utc' format='iso' value=2021-01-01 13:05:47.868>

            .. seealso::
                :ref:`ephemerides_sec`, :meth:`~nenupy.astro.target.Target.meridian_transit`

        """
        return self.meridian_transit(
            t_min=time,
            duration=TimeDelta(48*3600, format='sec'),
            precision=precision,
            fast_compute=fast_compute
        )[0]

    def previous_meridian_transit(self,
            time: Time = Time.now(),
            precision: TimeDelta = TimeDelta(5, format='sec'),
            fast_compute: bool = True
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` previous meridian transit time.
            This method returns the previous transit time found before ``time``.

            :param time:
                Relative time used to searching for the previous meridian transit.
                Default is current time.
            :type time:
                :class:`~astropy.time.Time`
            :param precision:
                Temporal precision of the returned meridian transit value.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`
            :param fast_compute:
                If set to ``True``, a fast approximation is used during
                the computation of Local Sidereal Time.
                Default is ``True``.
            :type fast_compute:
                `bool`

            :returns:
                Previous meridian transit time.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.previous_meridian_transit(
                        time=Time("2021-01-01 12:00:00")
                    )
                Time object: scale='utc' format='iso' value=2020-12-31 13:09:43.855>

            .. seealso::
                :ref:`ephemerides_sec`, :meth:`~nenupy.astro.target.Target.meridian_transit`

        """
        return self.meridian_transit(
            t_min=time - TimeDelta(48*3600, format='sec'),
            duration=TimeDelta(48*3600, format='sec'),
            precision=precision,
            fast_compute=fast_compute
        )[-1]

    def rise_time(self,
        t_min: Time = Time.now(),
        elevation: u.Quantity = 0*u.deg,
        duration: TimeDelta = TimeDelta(86400, format='sec'),
        precision: TimeDelta = TimeDelta(5, format='sec'),
        ):
        """ Computes the :class:`~nenupy.astro.target.Target` rise time(s) above ``elevation``.
            This method returns all the rise times found in the time
            window ranging from ``t_min`` to ``t_min + duration``.

            :param t_min:
                Starting time of the temporal window within which
                rise times are looked for.
                Default is current time.
            :type t_min:
                :class:`~astropy.time.Time`
            :param elevation:
                Elevation above which the rise time is computed.
                Default is ``0 deg``.
            :type elevation:
                :class:`~astropy.units.Quantity`
            :param duration:
                Width of the temporal window within which
                rise times are looked for.
                Default is ``1 day``.
            :type duration:
                :class:`~astropy.time.TimeDelta`
            :param precision:
                Temporal precision of the returned meridian transit values.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`

            :returns:
                Rise times above a given elevation.
                If no rise times are found (because the requested
                time window doesn't contain any) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> import astropy.units as u
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.rise_time(
                        t_min=Time("2021-01-01"),
                        elevation=0*u.deg,
                        duration=TimeDelta(86400*2, format="sec")
                    )
                <Time object: scale='utc' format='iso' value=['2021-01-01 02:28:51.926' '2021-01-02 02:24:56.599']>

            .. seealso::
                :ref:`ephemerides_sec`

        """

        def _find_elevation_rise_time(times):
            """ """
            altaz_coordinates = self._get_source_coordinates(time=times).transform_to(
                AltAz(
                    obstime=times,
                    location=self.observer
                )
            )
            elevations = altaz_coordinates.alt
            return np.where(
                (elevations[:, :-1] <= elevation) & (elevations[:, 1:] >= elevation)
            )
        return self._find_crossing_times(
            finding_function=_find_elevation_rise_time,
            t_min=t_min,
            duration=duration,
            precision=precision
        )

    def next_rise_time(self,
            time: Time = Time.now(),
            elevation: u.Quantity = 0*u.deg,
            precision: TimeDelta = TimeDelta(5, format='sec')
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` next rise time above ``elevation``.
            This method returns the next rise time found after ``time``.

            :param time:
                Relative time used to searching for the next rise time.
                Default is current time.
            :type time:
                :class:`~astropy.time.Time`
            :param elevation:
                Elevation above which the rise time is computed.
                Default is ``0 deg``.
            :type elevation:
                :class:`~astropy.units.Quantity`
            :param precision:
                Temporal precision of the returned rise time value.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`

            :returns:
                Next rise time.
                If no rise time is found (because the source does not
                cross the elevation) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> import astropy.units as u
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.next_rise_time(
                        time=Time("2021-01-01"),
                        elevation=40*u.deg,
                    )
                <Time object: scale='utc' format='iso' value=2021-01-01 08:20:16.447>

            .. seealso::
                :ref:`ephemerides_sec`, :meth:`~nenupy.astro.target.Target.rise_time`

        """
        try:
            return self.rise_time(
                t_min=time,
                elevation=elevation,
                duration=TimeDelta(48*3600, format='sec'),
                precision=precision
            )[0]
        except IndexError:
            return Time([], format="jd")

    def previous_rise_time(self,
            time: Time = Time.now(),
            elevation: u.Quantity = 0*u.deg,
            precision: TimeDelta = TimeDelta(5, format='sec')
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` previous rise time above ``elevation``.
            This method returns the previous rise time found after ``time``.

            :param time:
                Relative time used to searching for the previous rise time.
                Default is current time.
            :type time:
                :class:`~astropy.time.Time`
            :param elevation:
                Elevation above which the rise time is computed.
                Default is ``0 deg``.
            :type elevation:
                :class:`~astropy.units.Quantity`
            :param precision:
                Temporal precision of the returned rise time value.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`

            :returns:
                Previous rise time.
                If no rise time is found (because the source does not
                cross the elevation) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> import astropy.units as u
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.previous_rise_time(
                        time=Time("2021-01-01"),
                        elevation=40*u.deg,
                    )
                <Time object: scale='utc' format='iso' value=2020-12-31 08:24:12.434>

            .. seealso::
                :ref:`ephemerides_sec`, :meth:`~nenupy.astro.target.Target.rise_time`

        """
        try:
            return self.rise_time(
                t_min=time - TimeDelta(48*3600, format='sec'),
                elevation=elevation,
                duration=TimeDelta(48*3600, format='sec'),
                precision=precision
            )[-1]
        except IndexError:
            return Time([], format="jd")

    def set_time(self,
        t_min: Time = Time.now(),
        elevation: u.Quantity = 0*u.deg,
        duration: TimeDelta = TimeDelta(86400, format='sec'),
        precision: TimeDelta = TimeDelta(5, format='sec'),
        ):
        """ Computes the :class:`~nenupy.astro.target.Target` set time(s) below ``elevation``.
            This method returns all the set times found in the time
            window ranging from ``t_min`` to ``t_min + duration``.

            :param t_min:
                Starting time of the temporal window within which
                set times are looked for.
                Default is current time.
            :type t_min:
                :class:`~astropy.time.Time`
            :param elevation:
                Elevation below which the set time is computed.
                Default is ``0 deg``.
            :type elevation:
                :class:`~astropy.units.Quantity`
            :param duration:
                Width of the temporal window within which
                set times are looked for.
                Default is ``1 day``.
            :type duration:
                :class:`~astropy.time.TimeDelta`
            :param precision:
                Temporal precision of the returned meridian transit values.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`

            :returns:
                Set times below a given elevation.
                If no set times are found (because the requested
                time window doesn't contain any) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> import astropy.units as u
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.set_time(
                        t_min=Time("2021-01-01"),
                        elevation=0*u.deg,
                        duration=TimeDelta(86400*2, format="sec")
                    )
                <Time object: scale='utc' format='iso' value=['2021-01-01 23:42:41.174' '2021-01-02 23:38:45.188']>

            .. seealso::
                :ref:`ephemerides_sec`

        """

        def _find_elevation_set_time(times):
            """ """
            altaz_coordinates = self._get_source_coordinates(time=times).transform_to(
                AltAz(
                    obstime=times,
                    location=self.observer
                )
            )
            elevations = altaz_coordinates.alt
            return np.where(
                (elevations[:, :-1] >= elevation) & (elevations[:, 1:] <= elevation)
            )
        return self._find_crossing_times(
            finding_function=_find_elevation_set_time,
            t_min=t_min,
            duration=duration,
            precision=precision
        )

    def next_set_time(self,
            time: Time = Time.now(),
            elevation: u.Quantity = 0*u.deg,
            precision: TimeDelta = TimeDelta(5, format='sec')
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` next set time below ``elevation``.
            This method returns the next set time found after ``time``.

            :param time:
                Relative time used to searching for the next set time.
                Default is current time.
            :type time:
                :class:`~astropy.time.Time`
            :param elevation:
                Elevation below which the set time is computed.
                Default is ``0 deg``.
            :type elevation:
                :class:`~astropy.units.Quantity`
            :param precision:
                Temporal precision of the returned set time value.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`

            :returns:
                Next set time.
                If no set time is found (because the source does not
                cross the elevation) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> import astropy.units as u
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.next_set_time(
                        time=Time("2021-01-01"),
                        elevation=40*u.deg,
                    )
                <Time object: scale='utc' format='iso' value=2021-01-01 17:51:17.312>

            .. seealso::
                :ref:`ephemerides_sec`, :meth:`~nenupy.astro.target.Target.set_time`

        """
        try:
            return self.set_time(
                t_min=time,
                elevation=elevation,
                duration=TimeDelta(48*3600, format='sec'),
                precision=precision
            )[0]
        except IndexError:
            return Time([], format="jd")

    def previous_set_time(self,
            time: Time = Time.now(),
            elevation: u.Quantity = 0*u.deg,
            precision: TimeDelta = TimeDelta(5, format='sec')
        ) -> Time:
        """ Computes the :class:`~nenupy.astro.target.Target` previous set time below ``elevation``.
            This method returns the next set time found before ``time``.

            :param time:
                Relative time used to searching for the previous set time.
                Default is current time.
            :type time:
                :class:`~astropy.time.Time`
            :param elevation:
                Elevation below which the set time is computed.
                Default is ``0 deg``.
            :type elevation:
                :class:`~astropy.units.Quantity`
            :param precision:
                Temporal precision of the returned set time value.
                Default is ``5 sec``.
            :type precision:
                :class:`~astropy.time.TimeDelta`

            :returns:
                Previous set time.
                If no set time is found (because the source does not
                cross the elevation) an empty
                :class:`~astropy.time.Time` object is returned.
            :rtype:
                :class:`~astropy.time.Time`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> from astropy.time import Time, TimeDelta
                >>> import astropy.units as u
                >>> cyg_a = FixedTarget.from_name("Cyg A")
                >>> cyg_a.previous_set_time(
                        time=Time("2021-01-01"),
                        elevation=40*u.deg,
                    )
                <Time object: scale='utc' format='iso' value=2020-12-31 17:55:12.639>

            .. seealso::
                :ref:`ephemerides_sec`, :meth:`~nenupy.astro.target.Target.set_time`

        """
        try:
            return self.set_time(
                t_min=time - TimeDelta(48*3600, format='sec'),
                elevation=elevation,
                duration=TimeDelta(48*3600, format='sec'),
                precision=precision
            )[-1]
        except IndexError:
            return Time([], format="jd")

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    @abstractmethod
    def _get_source_coordinates(self, time: Time) -> SkyCoord:
        """ Abstract method that must be replaced by subclasses. """
        pass

    @staticmethod
    def _find_crossing_times(
            finding_function: Callable,
            t_min: Time,
            duration: TimeDelta,
            precision: TimeDelta
        ) -> np.ndarray:
        """ """
        # Set t_min to a higher dimension in preparation for multiple matches
        t_min = t_min.reshape((1,))

        # At each iteration, dt will be reduced by down_factor
        down_factor = 5
        dt = duration/down_factor

        # If duration is too big, intial dt is quite big as well
        # therefore, we set the max dt to 6h
        max_dt = TimeDelta(6*3600, format="sec")
        if dt > max_dt:
            down_factor = int(np.ceil(duration/max_dt))
            dt = duration/down_factor

        # Loop until the precision is reached
        while dt*down_factor > precision:

            # Prepare the time array upon which the coordinates are computed
            n_steps = np.ceil(duration/dt)
            times = t_min[:, None] + np.arange(n_steps + 1) * dt

            # Find the indices depending on the function to apply
            transit_indices = finding_function(times)

            # Update t_min at the spots where the transit(s) have been found
            t_min = times[transit_indices]

            if t_min.size == 0:
                # Nothing has been found
                return Time([], format='jd')
            elif t_min.isscalar:
                t_min = t_min.reshape((1,))

            # Next loop will occur on the last time step only
            duration = dt
            dt /= down_factor

        return times[transit_indices] + dt/2.
# ============================================================= #
# ============================================================= #

# ============================================================= #
# ------------------------ FixedTarget ------------------------ #
# ============================================================= #
class FixedTarget(Target):
    """ Class to handle astronomical targets outside the Solar System.

        .. versionadded:: 2.0.0

        :param coordinates:
        :type coordinates:
            :class:`~astropy.coordinates.SkyCoord`
        :param observer:
        :type observer:
            :class:`~astropy.coordinates.EarthLocation`
        :param time:
        :type time:
            :class:`~astropy.time.Time`

        .. rubric:: Attributes Summary

        .. autosummary::

            ~Target.coordinates
            ~Target.time
            ~Target.observer
            ~Target.is_circumpolar
            ~Target.culmination_azimuth

        .. rubric:: Methods Summary

        .. autosummary::

            ~FixedTarget.from_name
            ~Target.meridian_transit
            ~Target.next_meridian_transit
            ~Target.previous_meridian_transit
            ~Target.azimuth_transit
            ~Target.rise_time
            ~Target.next_rise_time
            ~Target.previous_rise_time
            ~Target.set_time
            ~Target.next_set_time
            ~Target.previous_set_time

        .. rubric:: Attributes and Methods Documentation

    """

    def __init__(self,
            coordinates: SkyCoord,
            time: Time = Time.now(),
            observer: EarthLocation = nenufar_position
        ):
        super().__init__(
            coordinates=coordinates,
            observer=observer,
            time=time
        )

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def horizontal_coordinates(self):
        """ """
        return super().horizontal_coordinates[:, 0]

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    @classmethod
    def from_name(cls,
            name: str,
            time: Time = Time.now(),
            observer: EarthLocation = nenufar_position
        ) -> FixedTarget:
        """ Instantiates a :class:`~nenupy.astro.target.FixedTarget`
            object from a name that could be resolved by `Simbad <http://simbad.u-strasbg.fr/simbad/>`_.

            :param name:
                Source name.
            :type name:
                `str`
            :param time:
                Time at which the source is looked at.
                Default is current time.
            :type time:
                :class:`~astropy.time.Time`
            :param observer:
                Earth location from where the source is observed.
                Default is NenuFAR's location.
            :type observer:
                :class:`~astropy.coordinates.EarthLocation`

            :returns:
                :class:`~nenupy.astro.target.FixedTarget` instance.
            :rtype:
                :class:`~nenupy.astro.target.FixedTarget`

            :Example:
                >>> from nenupy.astro.target import FixedTarget
                >>> cyg_a = FixedTarget.from_name("Cyg A")

        """

        if name.lower() in common_sources.keys():
            src = common_sources[name.lower()]
            source = SkyCoord(src["ra"], src["dec"], unit="deg")
        else:
            # Retrieve the Simbad coordinates
            source = SkyCoord.from_name(name)

        return cls(
            coordinates=source,
            observer=observer,
            time=time
        )

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    def _get_source_coordinates(self, time: Time):
        """ """
        return SkyCoord(
            ra=np.repeat(self.coordinates.ra.deg, time.size).reshape(time.shape),
            dec=np.repeat(self.coordinates.dec.deg, time.size).reshape(time.shape),
            unit="deg",
            frame=self.coordinates.frame
        )
        # return self.coordinates
# ============================================================= #
# ============================================================= #

# ============================================================= #
# --------------------- SolarSystemTarget --------------------- #
# ============================================================= #
class SolarSystemTarget(Target):
    """ Class to handle Solar System targets.

        .. versionadded:: 2.0.0

        .. rubric:: Attributes Summary

        .. autosummary::

            ~Target.coordinates
            ~Target.time
            ~Target.observer
            ~Target.is_circumpolar
            ~Target.culmination_azimuth

        .. rubric:: Methods Summary

        .. autosummary::

            ~SolarSystemTarget.from_name
            ~Target.meridian_transit
            ~Target.next_meridian_transit
            ~Target.previous_meridian_transit
            ~Target.azimuth_transit
            ~Target.rise_time
            ~Target.next_rise_time
            ~Target.previous_rise_time
            ~Target.set_time
            ~Target.next_set_time
            ~Target.previous_set_time

        .. rubric:: Attributes and Methods Documentation

    """

    def __init__(self,
            name: str,
            coordinates: SkyCoord,
            time: Time = Time.now(),
            observer: EarthLocation = nenufar_position
        ):
        super().__init__(
            coordinates=coordinates,
            observer=observer,
            time=time
        )
        self.name = name

    # --------------------------------------------------------- #
    # --------------------- Getter/Setter --------------------- #
    @property
    def horizontal_coordinates(self):
        """ """
        return super().horizontal_coordinates[np.identity(self.time.size, dtype=bool)]

    # --------------------------------------------------------- #
    # ------------------------ Methods ------------------------ #
    @classmethod
    def from_name(cls,
            name: str,
            time: Time = Time.now(),
            observer: EarthLocation = nenufar_position
        ) -> SolarSystemTarget:
        """ """

        # Get the ICRS instance of the Solar System object
        source = solar_system_source(
            name=name,
            time=time,
            observer=observer
        )

        return cls(
            name=name,
            coordinates=source,
            observer=observer,
            time=time
        )

    # --------------------------------------------------------- #
    # ----------------------- Internal ------------------------ #
    def _get_source_coordinates(self, time: Time):
        """ """
        return solar_system_source(
            name=self.name,
            time=time,
            observer=self.observer
        )
# ============================================================= #
# ============================================================= #

import numpy,import astropy,from astropy


# Copyright 2020 Google LLC
#
# 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.

# python3
"""A fork of PlaNet model.

Archive paper: https://arxiv.org/abs/1811.04551
OSS code repo: https://github.com/google-research/planet
"""
# pylint:disable=missing-docstring

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

import gin
import gym
import numpy as np
import tensorflow.compat.v1 as tf
import tensorflow.compat.v2.summary as tfs
from tensorflow_probability import distributions as tfd
from world_models.imported_models import reward_models
from world_models.utils import npz
from world_models.utils import visualization

from tensorflow.python.distribute import values

gin.external_configurable(tf.distribute.MirroredStrategy,
                          'tf.distribute.MirroredStrategy')

def static_scan(fn, inputs, start, reverse=False):
  # pylint: disable=expression-not-assigned
  # pylint: disable=cell-var-from-loop
  """drop-in replacement for tf.scan.

  tf.scan has some issues with multiple devices.
  """
  last = start
  outputs = [[] for _ in tf.nest.flatten(start)]
  indices = range(tf.nest.flatten(inputs)[0].shape[0])
  if reverse:
    indices = reversed(indices)
  for index in indices:
    inp = tf.nest.map_structure(lambda x: x[index], inputs)
    last = fn(last, inp)
    [o.append(l) for o, l in zip(outputs, tf.nest.flatten(last))]
  if reverse:
    outputs = [list(reversed(x)) for x in outputs]
  outputs = [tf.stack(x, 0) for x in outputs]
  return tf.nest.pack_sequence_as(start, outputs)

@gin.configurable
class RecurrentStateSpaceModel(object):

  def __init__(self,
      stoch_size=30,
      deter_size=200,
      min_stddev=0.1,
      layers=1,
      reward_layers=3,
      units=300,
      free_nats=3.0,
      reward_loss_multiplier=10,
      frame_size=(64, 64, 3),
      task=gin.REQUIRED,
      reward_from_frames=False,
      reward_stop_gradient=False,
      include_frames_in_prediction=False,
      activation=tf.nn.relu):
    self._action_space = task.create_env().action_space
    self._stoch_size = stoch_size
    self._deter_size = deter_size
    self._min_stddev = min_stddev
    self._num_layers = layers
    self._num_reward_layers = reward_layers
    self._num_units = units
    self._free_nats = free_nats
    self._include_frames_in_prediction = include_frames_in_prediction
    self._activation = activation
    self._cell = tf.keras.layers.GRUCell(self._deter_size)
    self._prior_tpl = tf.make_template('prior', self._prior)
    self._posterior_tpl = tf.make_template('posterior', self._posterior)
    self._encoder_tpl = tf.make_template('encoder', self._encoder)
    self._reward_loss_mul = reward_loss_multiplier
    self._frame_size = list(frame_size)
    self._reward_from_frames = reward_from_frames
    self._reward_stop_gradient = reward_stop_gradient
    self._predict_frames_tpl = tf.make_template(
        'predict_frames', self._predict_frames, out_shape=self._frame_size)
    self._predict_reward_tpl = tf.make_template(
        'predict_reward', self._predict_reward, out_shape=[1])

  @property
  def is_discrete_action(self):
    return isinstance(self._action_space, gym.spaces.Discrete)

  def get_trackables(self):
    return {
        'prior_tpl': self._prior_tpl,
        'posterior_tpl': self._posterior_tpl,
        'encoder_tpl': self._encoder_tpl,
        'predict_frames_tpl': self._predict_frames_tpl,
        'predict_reward_tpl': self._predict_reward_tpl,
        'cell': self._cell,
    }

  def initialize(self, batch_size):
    return {
        'mean':
          tf.zeros([batch_size, self._stoch_size]),
        'std':
          tf.zeros([batch_size, self._stoch_size]),
        'stoch':
          tf.zeros([batch_size, self._stoch_size]),
        'deter':
          self._cell.get_initial_state(
              batch_size=batch_size, dtype=tf.float32)
    }

  def compute_losses(self, obs):
    image = obs['image']
    action = obs['action']
    reward = obs['reward']
    state = self.initialize(tf.shape(image)[0])
    state['rewards'] = reward
    priors, posteriors = self.observe(action, image, state)
    features = self._get_features(posteriors)
    frames = self._predict_frames_tpl(features)
    if self._reward_from_frames:
      rewards = self._predict_reward_tpl(frames.mode(), reward[:, -1])
    else:
      rewards = self._predict_reward_tpl(features, reward[:, -1])
    obs_likelihood = frames.log_prob(image)
    reward_likelihood = rewards.log_prob(tf.to_float(reward))

    divergence = tfd.kl_divergence(
        self._get_distribution(posteriors), self._get_distribution(priors))
    divergence = tf.maximum(self._free_nats, divergence)
    loss = tf.reduce_mean(divergence - obs_likelihood -
                          reward_likelihood * self._reward_loss_mul)
    frames_mode = tf.clip_by_value((frames.mode() + 0.5) * 255, 0, 255)
    return (loss, tf.reduce_mean(reward_likelihood), tf.reduce_mean(divergence),
            tf.cast(frames_mode, dtype=tf.uint8), rewards.mode(), obs['reward'],
            tf.reduce_mean(tf.math.squared_difference(frames_mode, image)))

  def observe(self, actions, images, state):
    embedded_obs = self._encoder_tpl(images)
    if self.is_discrete_action:
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.to_float(actions)
    actions = tf.transpose(actions, [1, 0, 2])
    embedded_obs = tf.transpose(embedded_obs, [1, 0, 2])
    state.pop('rewards')
    priors, posteriors = static_scan(
        lambda prev, inp: self._posterior_tpl(prev[1], *inp),
        (actions, embedded_obs), (state, state))
    priors = {
        key: tf.transpose(value, [1, 0, 2]) for key, value in priors.items()
    }
    posteriors = {
        key: tf.transpose(value, [1, 0, 2])
        for key, value in posteriors.items()
    }
    return priors, posteriors

  def predict(self, actions, state):
    if isinstance(self._action_space, gym.spaces.Discrete):
      actions = tf.one_hot(
          actions[:, :, 0], self._action_space.n, dtype=tf.float32)
    else:
      actions = tf.to_float(actions)
    actions = tf.transpose(actions, [1, 0, 2])
    rewards = tf.to_float(state.pop('rewards'))
    priors = static_scan(self._prior_tpl, actions, state)
    priors = {
        key: tf.transpose(value, [1, 0, 2]) for key, value in priors.items()
    }
    features = self._get_features(priors)
    results = {}
    if self._reward_from_frames:
      frames = self._predict_frames_tpl(features).mode()
      results['reward'] = self._predict_reward_tpl(frames, rewards).mode()
    else:
      results['reward'] = self._predict_reward_tpl(features, rewards).mode()
    if self._include_frames_in_prediction:
      results['image'] = tf.cast(
          tf.clip_by_value(
              (self._predict_frames_tpl(features).mode() + 0.5) * 255, 0, 255),
          dtype=tf.uint8)
    return results

  def _get_distribution(self, states):
    return tfd.MultivariateNormalDiag(states['mean'], states['std'])

  def _get_features(self, states):
    return tf.concat([states['stoch'], states['deter']], -1)

  def _prior(self, prev_state, prev_action):
    hidden = tf.concat([prev_state['stoch'], prev_action], -1)
    for _ in range(self._num_layers):
      hidden = tf.layers.dense(hidden, self._num_units,
                               self._activation)
    hidden, deter = self._cell(hidden, [prev_state['deter']])
    deter = deter[0]
    for _ in range(self._num_layers):
      hidden = tf.layers.dense(hidden, self._num_units,
                               self._activation)
    mean, std = tf.split(
        tf.layers.dense(hidden, 2 * self._stoch_size), 2, -1)
    std = tf.nn.softplus(std) + self._min_stddev
    stoch = tfd.MultivariateNormalDiag(mean, std).sample()
    return {'mean': mean, 'std': std, 'stoch': stoch, 'deter': deter}

  def _posterior(self, prev_state, prev_action, embedded_obs):
    prior = self._prior_tpl(prev_state, prev_action)
    hidden = tf.concat([prior['deter'], embedded_obs], -1)
    for _ in range(self._num_layers):
      hidden = tf.layers.dense(hidden, self._num_units,
                               self._activation)
    mean, std = tf.split(
        tf.layers.dense(hidden, 2 * self._stoch_size), 2, -1)
    std = tf.nn.softplus(std) + self._min_stddev
    stoch = tfd.MultivariateNormalDiag(mean, std).sample()
    post = {'mean': mean, 'std': std, 'stoch': stoch, 'deter': prior['deter']}
    return prior, post

  def _encoder(self, images):
    kwargs = dict(strides=2, activation=tf.nn.relu)
    images = tf.to_float(images)
    hidden = tf.reshape(images, [-1] + images.shape[2:].as_list())
    hidden = tf.layers.conv2d(hidden, 32, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 64, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 128, 4, **kwargs)
    hidden = tf.layers.conv2d(hidden, 256, 4, **kwargs)
    hidden = tf.layers.flatten(hidden)
    assert hidden.shape[1:].as_list() == [1024], hidden.shape.as_list()
    embedded_obs = tf.reshape(hidden, [
        tf.shape(images)[0],
        tf.shape(images)[1],
        np.prod(hidden.shape[1:].as_list())
    ])
    return embedded_obs

  def _predict_frames(self, features, out_shape):
    kwargs = dict(strides=2, activation=tf.nn.relu)
    hidden = tf.layers.dense(features, 1024, None)
    hidden = tf.reshape(hidden, [-1, 1, 1, hidden.shape[-1]])
    hidden = tf.layers.conv2d_transpose(hidden, 128, 5, **kwargs)
    hidden = tf.layers.conv2d_transpose(hidden, 64, 5, **kwargs)
    hidden = tf.layers.conv2d_transpose(hidden, 32, 6, **kwargs)
    mean = tf.layers.conv2d_transpose(hidden, 3, 6, strides=2)
    assert mean.shape[1:].as_list() == [64, 64, 3], mean.shape
    mean = tf.reshape(mean, tf.concat([tf.shape(features)[:-1], out_shape], 0))
    mean = tf.cast(mean, tf.float32)
    return tfd.Independent(tfd.Normal(mean, 1), len(out_shape))

  def _predict_reward(self, features, rewards, out_shape):
    if self._reward_stop_gradient:
      hidden = tf.stop_gradient(features)
    else:
      hidden = features
    if self._reward_from_frames:
      split_frames = [
          tf.squeeze(f, axis=1)
          for f in tf.split(features, features.shape[1], axis=1)
      ]
      mean = reward_models.reward_prediction_video_conv(split_frames, rewards,
                                                        len(split_frames))
      return tfd.Independent(tfd.Normal(mean, 1), len(out_shape))
    else:
      for _ in range(self._num_reward_layers):
        hidden = tf.layers.dense(hidden, self._num_units,
                                 self._activation)
      mean = tf.layers.dense(hidden, int(np.prod(out_shape)))
      mean = tf.reshape(mean, tf.concat([tf.shape(features)[:-1], out_shape],
                                        0))
      return tfd.Independent(tfd.Normal(mean, 1), len(out_shape))

@gin.configurable
def create_planet_reset_fn(model):
  @tf.function
  def reset(**kwargs):
    batch_size = kwargs['proposals']
    state = model.initialize(batch_size)
    state['rewards'] = tf.zeros([batch_size, 1])
    return state

  return reset

@gin.configurable
def create_planet_observe_fn(model, model_dir, strategy):
  with strategy.scope():
    checkpoint = tf.train.Checkpoint(**model.get_trackables())
    manager = tf.train.CheckpointManager(checkpoint, model_dir, max_to_keep=1)
    checkpoint.restore(manager.latest_checkpoint).expect_partial()

  @tf.function
  def observe(images, actions, rewards, state):
    images = tf.to_float(images) / 255.0 - 0.5
    # break down the inputs along the batch dimension to form equal sized
    # tensors in each replica.
    num_replicas = strategy.num_replicas_in_sync
    images = tf.split(images, num_replicas)
    actions = tf.split(actions, num_replicas)
    state = {key: tf.split(value, num_replicas) for key, value in state.items()}
    devices = values.ReplicaDeviceMap(strategy.extended.worker_devices)
    dist_images = values.PerReplica(devices, tuple(images))
    dist_actions = values.PerReplica(devices, tuple(actions))
    dist_state = []
    for i in range(num_replicas):
      dist_state.append({key: value[i] for key, value in state.items()})
    dist_state = values.PerReplica(devices, tuple(dist_state))
    _, dist_posteriors = strategy.experimental_run_v2(
        model.observe, args=(dist_actions, dist_images, dist_state))
    dist_posteriors = {
        key: strategy.experimental_local_results(value)
        for key, value in dist_posteriors.items()
    }
    posteriors = {
        key: tf.concat(value, axis=0) for key, value in dist_posteriors.items()
    }
    posteriors = {key: value[:, -1] for key, value in posteriors.items()}
    posteriors['rewards'] = rewards[:, -1]
    return posteriors

  return observe

@gin.configurable
def create_planet_predict_fn(model, strategy):
  @tf.function
  def predict(actions, state):
    state = state.copy()
    # break down the inputs along the batch dimension to form equal sized
    # tensors in each replica.
    num_replicas = strategy.num_replicas_in_sync
    actions = tf.split(actions, num_replicas)
    state = {key: tf.split(value, num_replicas) for key, value in state.items()}
    devices = values.ReplicaDeviceMap(strategy.extended.worker_devices)
    dist_actions = values.PerReplica(devices, tuple(actions))
    dist_state = []
    for i in range(num_replicas):
      dist_state.append({key: value[i] for key, value in state.items()})
    dist_state = values.PerReplica(devices, tuple(dist_state))

    dist_predictions = strategy.experimental_run_v2(
        model.predict, args=(dist_actions, dist_state))
    dist_predictions = {
        key: strategy.experimental_local_results(value)
        for key, value in dist_predictions.items()
    }
    predictions = {
        key: tf.concat(value, axis=0)
        for key, value in dist_predictions.items()
    }
    return predictions

  return predict

@gin.configurable
def create_planet_train_fn(model: RecurrentStateSpaceModel = gin.REQUIRED,
    train_steps: int = gin.REQUIRED,
    batch: int = gin.REQUIRED,
    duration: int = gin.REQUIRED,
    learning_rate: float = gin.REQUIRED,
    model_dir=gin.REQUIRED,
    strategy: tf.distribute.Strategy = gin.REQUIRED,
    save_rewards: bool = True):
  """creates a train_fn to train the `tf.Estimator` referenced in state.

  Args:
    model: a reference to the model.
    train_steps: number of training steps.
    batch: the batch size.
    duration: how many timesteps to include in a single video sequence.
    learning_rate: learning rate.
    model_dir: the path to model directory.
    strategy: a tf.distribute.Strategy object.
    save_rewards: whether or not to save the predicted rewards.

  Returns:
    A train_fn with the following positional arguments:
        * data_path: the path to all episodes.
      This function returns nothing.
  """
  iterator = None
  optimizer = tf.keras.optimizers.Adam(learning_rate, epsilon=1e-3)

  @tf.function
  def train_step(obs):

    def train_iter(obs):
      obs = obs.copy()
      obs['image'] = tf.to_float(obs['image']) / 255.0 - 0.5
      with tf.GradientTape() as tape:
        output = model.compute_losses(obs)
        loss, reward_loss, divergence, frames = output[:4]
        pred_rewards, true_rewards, frame_loss = output[4:]
      variables = tape.watched_variables()
      grads = tape.gradient(loss, variables)
      grads, _ = tf.clip_by_global_norm(grads, 1000)
      optimizer.apply_gradients(zip(grads, variables))
      return loss, reward_loss, divergence, frames, pred_rewards, true_rewards, frame_loss

    return strategy.experimental_run_v2(train_iter, args=(obs,))

  def train_fn(data_path):
    """A train_fn to train the planet model."""
    nonlocal iterator
    nonlocal optimizer

    with strategy.scope():
      global_step = tf.Variable(0, dtype=tf.int64, trainable=False)
      checkpoint = tf.train.Checkpoint(
          global_step=global_step,
          optimizer=optimizer,
          **model.get_trackables())
      manager = tf.train.CheckpointManager(checkpoint, model_dir, max_to_keep=1)
      checkpoint.restore(manager.latest_checkpoint)
    if iterator is None:
      dataset = npz.load_dataset_from_directory(data_path, duration, batch)
      dataset = strategy.experimental_distribute_dataset(dataset)
      iterator = dataset

    writer = tfs.create_file_writer(model_dir)
    tfs.experimental.set_step(global_step)
    true_rewards, pred_rewards = None, None
    with writer.as_default():
      for step, obs in enumerate(iterator):
        if step > train_steps:
          if save_rewards:
            # We are only saving the last training batch.
            reward_dir = os.path.join(model_dir, 'train_rewards')
            true_rewards = strategy.experimental_local_results(true_rewards)
            pred_reward = strategy.experimental_local_results(pred_rewards)
            true_rewards = np.concatenate([x.numpy() for x in true_rewards])
            pred_reward = np.concatenate([x.numpy() for x in pred_reward])
            rewards_to_save = {'true': true_rewards, 'pred': pred_reward}
            npz.save_dictionary(rewards_to_save, reward_dir)
          break
        (loss, reward_loss, divergence, frames, pred_rewards, true_rewards,
         frame_loss) = train_step(obs)
        if step % 100 == 0:
          loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, loss)
          reward_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN,
                                        reward_loss)
          divergence = strategy.reduce(tf.distribute.ReduceOp.MEAN, divergence)
          frame_loss = strategy.reduce(tf.distribute.ReduceOp.MEAN, frame_loss)
          frames = strategy.experimental_local_results(frames)
          frames = tf.concat(frames, axis=0)
          pred_reward = strategy.experimental_local_results(pred_rewards)
          pred_reward = tf.concat(pred_reward, axis=0)
          tf.logging.info('loss at step %d: %f', step, loss)
          tfs.scalar('loss/total', loss)
          tfs.scalar('loss/reward', reward_loss)
          tfs.scalar('loss/divergence', divergence)
          tfs.scalar('loss/frames', frame_loss)
          tfs.experimental.write_raw_pb(
              visualization.py_gif_summary(
                  tag='predictions/frames',
                  images=frames.numpy(),
                  max_outputs=6,
                  fps=20))
          ground_truth_rewards = (
              tf.concat(
                  strategy.experimental_local_results(obs['reward']),
                  axis=0)[:, :, 0])
          rewards = pred_reward[:, :, 0]
          signals = tf.stack([ground_truth_rewards, rewards], axis=1)
          visualization.py_plot_1d_signal(
              name='predictions/reward',
              signals=signals.numpy(),
              labels=['ground_truth', 'prediction'],
              max_outputs=6)
        global_step.assign_add(1)

      manager.save(global_step)

  return train_fn


#M3 -- Meka Robotics Robot Components
#Copyright (c) 2010 Meka Robotics
#Author: edsinger@mekabot.com (Aaron Edsinger)

#M3 is free software: you can redistribute it and/or modify
#it under the terms of the GNU Lesser General Public License as published by
#the Free Software Foundation, either version 3 of the License, or
#(at your option) any later version.

#M3 is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU Lesser General Public License for more details.

#You should have received a copy of the GNU Lesser General Public License
#along with M3.  If not, see <http://www.gnu.org/licenses/>.

from m3.chain import M3Chain
import m3.toolbox as m3t
import m3.toolbox_ctrl as m3tc
import scipy.linalg
import numpy as nu
from m3.unit_conversion import *

#Wrapper only
class M3Arm(M3Chain):
	"""Wrapper for 7DOF SEA Arm"""
	def __init__(self,name):
		M3Chain.__init__(self,name,ndof=7,ctype='m3arm')
		self.set_fts_transform(nu.identity(4,float)) #default
		self.ikjt_q=None
	# #############################################################################

	"""Force-torque sensor support.
	Assumes that wrench is 6x1 Numeric array.
	See documentation for frame definitions"""

	def set_fts_transform(self,T):
		""" Set the homogenous transform from the force-torque-sensor frame to eff frame"""
		self.S2E = nu.array(T,float) #Transform point in sensor frame to eff
		self.E2S =  scipy.linalg.inv(self.S2E)  #Transform point in eff to sensor frame.
		self.FS2FE = m3tc.force_moment_transform(self.S2E)  #Transform wrench in sensor frame to eff frame
		self.FE2FS = m3tc.force_moment_transform(self.E2S)  #Transform wrench in eff frame to sensor frame

	def fts_wrench_to_tool_wrench(self,wrench):
		effw=nu.matrixmultiply(self.FS2FE,wrench)
		endw=nu.matrixmultiply(self.eff_wrench_2_end_wrench_transform(),effw)
		return self.end_wrench_2_tool_wrench(endw)

	def tool_wrench_to_fts_wrench(self,wrench):
		endw=self.tool_wrench_2_end_wrench(wrench)
		effw=nu.matrixmultiply(self.end_wrench_2_eff_wrench_transform(),endw)
		return nu.matrixmultiply(self.FE2FS,effw)

	def fts_wrench_to_joint_torques(self,wrench):
		effw=nu.matrixmultiply(self.FS2FE,wrench)
		endw=nu.matrixmultiply(self.eff_wrench_2_end_wrench_transform(),effw)
		return self.end_wrench_2_joint_torques(endw)

	def joint_torques_to_fts_wrench(self,tq):
		endw=self.joint_torques_2_end_wrench(tq)
		effw=nu.matrixmultiply(self.end_wrench_2_eff_wrench_transform(),endw)
		return nu.matrixmultiply(self.FE2FS,effw)

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

	def start_ikjt_simple(self):
		self.ikjt_q=self.get_theta_rad()

	def step_ikjt_simple(self,target,step_size):
		""" A simple method for doing inverse-kinematics using the Jacobian Transpose. Because
	of the redundant DOF , this only good for small excursions from a known
	good posture in joint-space. Othewise strange elbow configurations can arise.
	"""
		self.set_virtual_theta_rad(self.ikjt_q)
		x=self.get_virtual_end_position()
		dx=(target-x)*step_size
		verror=nu.sqrt(sum((target-x)**2))
		dx.resize(6) #roll/pitch/yaw=0
		dq= nu.matrixmultiply(nu.transpose(self.vJ),dx)
		self.ikjt_q=self.ikjt_q+dq
		return verror,self.ikjt_q