##############################################################################
#
# Copyright (c) 2003-2018 by The University of Queensland
# http://www.uq.edu.au
#
# Primary Business: Queensland, Australia
# Licensed under the Apache License, version 2.0
# http://www.apache.org/licenses/LICENSE-2.0
#
# Development until 2012 by Earth Systems Science Computational Center (ESSCC)
# Development 2012-2013 by School of Earth Sciences
# Development from 2014 by Centre for Geoscience Computing (GeoComp)
#
##############################################################################

"""Base classes for forward models"""

from __future__ import division, print_function

__copyright__="""Copyright (c) 2003-2018 by The University of Queensland
http://www.uq.edu.au
Primary Business: Queensland, Australia"""
__license__="""Licensed under the Apache License, version 2.0
http://www.apache.org/licenses/LICENSE-2.0"""
__url__="https://launchpad.net/escript-finley"

__all__ = ['ForwardModel','ForwardModelWithPotential']

from esys.downunder.coordinates import makeTransformation
from esys.escript.linearPDEs import LinearSinglePDE
from esys.escript.util import *
import numpy as np

class ForwardModel(object):
    """
    An abstract forward model that can be plugged into a cost function.
    Subclasses need to implement `getDefect()`, `getGradient()`, and possibly
    `getArguments()` and 'getCoordinateTransformation'.
    """
    def __init__(self):
        pass

    def getArguments(self, x):
        return ()

    def getCoordinateTransformation(self):
        return None

    def getDefect(self, x, *args):
        raise NotImplementedError

    def getGradient(self, x, *args):
        raise NotImplementedError

class ForwardModelWithPotential(ForwardModel):
    """
    Base class for a forward model using a potential such as magnetic or
    gravity. It defines a cost function:

        defect = 1/2 sum_s integrate( ( weight_i[s] * ( r_i - data_i[s] ) )**2 )

    where s runs over the survey, weight_i are weighting factors, data_i are
    the data, and r_i are the results produced by the forward model.
    It is assumed that the forward model is produced through postprocessing
    of the solution of a potential PDE.
    """
    def __init__(self, domain, w, data,  coordinates=None,
                                 fixPotentialAtBottom=False,
                                 tol=1e-8):
        """
        initializes a new forward model with potential.

        :param domain: domain of the model
        :type domain: `Domain`
        :param w: data weighting factors
        :type w: ``Vector`` or list of ``Vector``
        :param data: data
        :type data: ``Vector`` or list of ``Vector``
        :param coordinates: defines coordinate system to be used
        :type coordinates: `ReferenceSystem` or `SpatialCoordinateTransformation`
        :param fixPotentialAtBottom: if true potential is fixed to zero at the bottom of the domain
                                     in addition to the top.
        :type fixPotentialAtBottom: ``bool``
        :param tol: tolerance of underlying PDE
        :type tol: positive ``float``
        """
        super(ForwardModelWithPotential, self).__init__()
        self.__domain = domain
        self.__trafo = makeTransformation(domain, coordinates)

        try:
            n=len(w)
            m=len(data)
            if not m == n:
                raise ValueError("Length of weight and data must be the same.")
            self.__weight = w
            self.__data = data
        except TypeError:
            self.__weight = [w]
            self.__data = [data]

        BX = boundingBox(domain)
        DIM = domain.getDim()
        x = domain.getX()
        self.__pde=LinearSinglePDE(domain)
        self.__pde.getSolverOptions().setTolerance(tol)
        self.__pde.setSymmetryOn()
        z=x[DIM-1]
        q0=whereZero(z-BX[DIM-1][1])
        if fixPotentialAtBottom: q0+=whereZero(z-BX[DIM-1][0])
        self.__pde.setValue(q=q0)

        self.edge_lengths=np.asarray(boundingBoxEdgeLengths(domain))
        self.diameter=1./sqrt(sum(1./self.edge_lengths**2))

        self.__origweight=[]
        for s in range(len(self.__weight)):
            # save a copy of the original weights in case of rescaling
            self.__origweight.append(1.*self.__weight[s])

        if not self.__trafo.isCartesian():
            fd=1./self.__trafo.getScalingFactors()
            fw=self.__trafo.getScalingFactors()*sqrt(self.__trafo.getVolumeFactor())
            for s in range(len(self.__weight)):
                self.__weight[s] = fw * self.__weight[s]
                self.__data[s]   = fd * self.__data[s]

    def _rescaleWeights(self, scale=1., fetch_factor=1.):
        """
        rescales the weights such that

        *sum_s integrate( ( weight_i[s] *data_i[s]) (weight_j[s]*1/L_j) * L**2 * fetch_factor )=scale*
        """
        if not scale > 0:
             raise ValueError("Value for scale must be positive.")
        A=0
        # copy back original weights before rescaling
        self.__weight=[1.*ow for ow in self.__origweight]

        for s in range(len(self.__weight)):
            if self.__data[s].getShape() == ():
               ff=self.__weight[s]**2*self.__data[s]/length(self.edge_lengths)
            else:
               ff=inner(self.__weight[s], self.__data[s]) * inner(self.__weight[s], 1/self.edge_lengths)
            A += integrate(abs(ff * fetch_factor))
        if A > 0:
            A=sqrt(scale/A)/self.diameter
            if not self.__trafo.isCartesian():
                A*=self.__trafo.getScalingFactors()*sqrt(self.__trafo.getVolumeFactor())
            for s in range(len(self.__weight)):
                self.__weight[s]*=A
        else:
            raise ValueError("Rescaling of weights failed.")

    def getDomain(self):
        """
        Returns the domain of the forward model.

        :rtype: `Domain`
        """
        return self.__domain

    def getMisfitWeights(self):
        """
        Returns the weights of the misfit function

        :rtype: ``list`` of ``Data``
        """
        return self.__weight

    def getData(self):
        """
        Returns the data

        :rtype: ``list`` of ``Data``
        """
        return self.__data
    def getDataFunctionSpace(self):
        """
        Returns the ``FunctionSpace`` of the data

        :rtype: ``FunctionSpace``
        """
        return self.getData()[0].getFunctionSpace()

    def getCoordinateTransformation(self):
        """
        returns the coordinate transformation being used

        :rtype: ``CoordinateTransformation``
        """
        return self.__trafo

    def getPDE(self):
        """
        Return the underlying PDE.

        :rtype: `LinearPDE`
        """
        return self.__pde

    def _getDefect(self, result):
        """
        Returns the defect value.

        :param result: a result vector
        :type result: `Vector`
        :rtype: ``float``
        """
        A=0.
        for s in range(len(self.__weight)):
            A += integrate( inner(self.__weight[s], self.__data[s]-result)**2 )
        return A/2

    def getDefectGradient(self, result):
        Y=0.
        for s in range(len(self.__weight)):
            Y = inner(self.__weight[s], self.__data[s]-result) * self.__weight[s] + Y
        return Y

    def getSurvey(self, index=None):
        """
        Returns the pair (data_index, weight_index), where data_i is the data
        of survey i, weight_i is the weighting factor for survey i.
        If index is None, all surveys will be returned in a pair of lists.
        """
        if index is None:
            return self.__data, self.__weight
        if index>=len(self.__data):
            raise IndexError("Forward model only has %d surveys"%len(self.__data))
        return self.__data[index], self.__weight[index]


import numpy as np
import pymc3 as pm
import pandas as pd
#import pickle
#import wandb
#import yaml
import warnings
from sklearn.cluster import k_means
from scipy.special import softmax, logsumexp, loggamma
from sklearn.metrics.pairwise import cosine_similarity
from .constants import *
import pkg_resources

# constants
#C=32
#M=3
#P=2

def dirichlet(node_name, a, shape, scale=1, testval = None):
    # dirichlet reparameterized here because of stickbreaking bug
    # https://github.com/pymc-devs/pymc3/issues/4733
    X = pm.Gamma(f'gamma_{node_name}', mu = a, sigma = scale, shape = shape, testval = testval)
    Y = pm.Deterministic(node_name, (X/X.sum(axis = (X.ndim-1))[...,None]))
    return Y

def load_config(config_fp):

    # load the yaml file
    with open(config_fp, 'r') as f:
        config = yaml.safe_load(f)
        print(f"Loaded configuration file {config_fp}")

    # remove any parameters not applicable to selected data source
    ds = config.pop('dataset')
    ds[ds['dataset_sel']].update({'dataset_sel': ds['dataset_sel']})

    # update dataset args to subsetted list
    config.update({'dataset': ds[ds['dataset_sel']]})

    ## handle seeding
    #config['dataset'].update({'data_rng': np.random.default_rng(config['dataset']['data_seed'])})
    #config['model'].update({'model_rng': np.random.default_rng(config['model']['model_seed'])})

    return config['dataset'], config['model'], config['pymc3']

def detect_naming_style(fp):

    # first column must have type at minimum
    df = pd.read_csv(fp, index_col=0, sep = None, engine = 'python')
    naming_style = 'unrecognized'

    # check if index is type style
    if df.index.isin(mut96).any(): naming_style = 'type'

    # check if index is type/subtype style
    else:
        df = df.reset_index()
        df = df.set_index(list(df.columns[0:2]))
        if df.index.isin(idx96).any(): naming_style = 'type/subtype'

    assert naming_style == 'type' or naming_style == 'type/subtype', \
            'Mutation type naming style could not be identified.\n'\
            '\tExpected either two column type/subtype (ex. C>A,ACA) or\n'\
            '\tsingle column type (ex A[C>A]A). See examples at COSMIC database.'

    return naming_style

def load_sigs(fp):
    warnings.warn("load_sigs is deprecated, see Damuta class", DeprecationWarning)

    naming_style = detect_naming_style(fp)

    if naming_style == 'type':
        # read in sigs
        sigs = pd.read_csv(fp, index_col = 0, sep = None, engine = 'python').reindex(mut96)
        # sanity check for matching mut96, should have no NA
        sel = (~sigs.isnull()).all(axis = 1)
        assert sel.all(), f'invalid signature definitions: null entry for types {list(sigs.index[~sel])}'
        # convert to pcawg convention
        sigs = sigs.set_index(idx96)

    elif naming_style == 'type/subtype':
        # read in sigs
        sigs = pd.read_csv(fp, index_col = (0,1), header=0).reindex(idx96)
        # sanity check for idx, should have no NA
        sel = (~sigs.isnull()).all(axis = 1)
        assert sel.all(), f'invalid signature definitions: null entry for types {list(sigs.index[~sel])}'

    # check colsums are 1
    sel = np.isclose(sigs.sum(axis=0), 1)
    assert sel.all(), f'invalid signature definitions: does not sum to 1 in columns {list(sigs.columns[~sel])}'

    assert all(sigs.index == idx96) or all(sigs.index == mut96), 'signature defintions failed to be read correctly'

    # force Jx96 and mut96 convention
    sigs = sigs.T
    sigs.columns = mut96

    return sigs

def load_counts(counts_fp):
    warnings.warn("load_counts is deprecated, see Damuta class", DeprecationWarning)
    counts = pd.read_csv(counts_fp, index_col = 0, header = 0)[mut96]
    assert counts.ndim == 2, 'Mutation counts failed to load. Check column names are mutation type (ex. A[C>A]A). See COSMIC database for more.'
    assert counts.shape[1] == 96, f'Expected 96 mutation types, got {counts.shape[1]}'
    return counts

def subset_samples(dataset, annotation, annotation_subset, sel_idx = 0):
    # subset sample ids by matching to annotation_subset
    # expect annotation_subset to be pd dataframe with ids as index

    if annotation_subset is None:
        return dataset, annotation

    # stop string being auto cast to list
    if type(annotation_subset) == str:
        annotation_subset = [annotation_subset]

    if annotation.ndim > 2:
        warnings.warn(f"More than one annotation is available per sample, selection index {sel_idx}", UserWarning)

    # annotation ids should match sample ids
    assert dataset.index.isin(annotation.index).any(), 'No sample ID matches found in dataset for the provided annotation'

    # reoder annotation (with gaps) to match dataset
    annotation = annotation.reindex(dataset.index)

    # partial matches allowed
    sel = np.fromiter((map(any, zip(*[annotation[annotation.columns[sel_idx]].str.contains(x) for x in annotation_subset] ))), dtype = bool)

    # type should appear in the type column of the lookup
    assert sel.any(), 'Cancer type subsetting yielded no selection. Check keywords?'

    dataset = dataset.loc[annotation.index[sel]]
    annotation = annotation.loc[annotation.index[sel]]
    return dataset, annotation

def save_checkpoint(fp, model, trace, dataset_args, model_args, pymc3_args, run_id):
    with open(f'{fp}', 'wb') as buff:
        pickle.dump({'model': model, 'trace': trace, 'dataset_args': dataset_args,
                     'model_args': model_args, 'pymc3_args': pymc3_args, 'run_id': run_id}, buff)
    print(f'checkpoint saved to {fp}')

def load_checkpoint(fn):
    with open(fn, 'rb') as buff:
        data = pickle.load(buff)
        print(f'checkpoint loaded from {fn}')
        wandb.init(id=data['run_id'], resume='allow')
        return data['model'], data['trace'], data['dataset_args'], data['model_args'], data['pymc3_args'], data['run_id']

def load_dataset(dataset_sel, counts_fp=None, annotation_fp=None, annotation_subset=None, seed=None,
                 data_seed = None, sig_defs_fp=None, sim_S=None, sim_N=None, sim_I=None, sim_tau_hyperprior=None,
                 sim_J=None, sim_K=None, sim_alpha_bias=None, sim_psi_bias=None, sim_gamma_bias=None, sim_beta_bias=None):
    # load counts, or simulated data - as specified by dataset_sel
    # seed -> rng as per https://albertcthomas.github.io/good-practices-random-number-generators/

    if dataset_sel == 'load_counts':
        dataset = load_counts(counts_fp)
        annotation = pd.read_csv(annotation_fp, index_col = 0, header = 0)
        dataset, annotation = subset_samples(dataset, annotation, annotation_subset)
        return dataset, annotation

    elif dataset_sel == 'sim_from_sigs':
        sig_defs = load_sigs(sig_defs_fp)
        dataset, sim_params = sim_from_sigs(sig_defs, sim_tau_hyperprior, sim_S, sim_N, sim_I, seed)
        return dataset, sim_params

    elif dataset_sel == 'sim_parametric':
        dataset, sim_params = sim_parametric(sim_J,sim_K,sim_S,sim_N,sim_alpha_bias,sim_psi_bias,sim_gamma_bias,sim_beta_bias,seed)
        return dataset, sim_params

    else:
        assert False, 'dataset selection not recognized'

def load_datasets(dataset_args):
    yargs = dataset_args.copy()
    ca = [load_dataset(counts_fp = j[0], annotation_fp = j[1], **yargs) for j in zip(yargs.pop('counts_fp'), yargs.pop('annotation_fp'))]
    counts = pd.concat([c[0] for c in ca ])
    annotation = pd.concat([a[1] for a in ca])
    return counts, annotation

def split_by_count(data, fraction, rng):
    # assumes data is a pandas df

    frac1 = data.apply(lambda r: rng.choice( np.repeat(np.arange(96), r), int(fraction*r.sum()), replace = False) , axis = 1)
    frac1 = pd.DataFrame(frac1.apply(lambda r: np.histogram(r, bins=96, range=(0, 96))[0]).to_list(),
                         index = data.index, columns = data.columns)
    #frac1 = (data * fraction).astype(int)
    #frac1 = frac1.apply(lambda r: np.histogram(r, bins=96, range=(0, 96))[0], axis = 1, result_type='expand')
    #frac1.columns = data.columns
    frac2 = data - frac1
    assert all(frac2 >= 0) and all(frac1 >= 0)
    assert np.all(data == frac1 + frac2), 'Splitting failed'
    return frac1, frac2

def split_by_S(data, fraction, rng):
    # assumes data is a pandas df with an index
    frac1 = data.sample(frac=fraction, random_state=np.random.RandomState(rng.bit_generator))
    frac2 = data.drop(frac1.index)
    return frac1, frac2

def split_data(counts, S_frac = 0.9, c_frac = 0.8, rng=np.random.default_rng()):
    # get train/val/test split
    trn, tst = split_by_S(counts, S_frac, rng)
    trn, val = split_by_count(trn, c_frac, rng)
    tst1, tst2 = split_by_count(tst, c_frac, rng)
    return trn, val, tst1, tst2

def get_tau(phi, eta):
    assert len(phi.shape) == 2 and len(eta.shape) == 3
    tau =  np.einsum('jpc,kpm->jkpmc', phi.reshape((-1,2,16)), eta).reshape((-1,96))
    return tau

def get_phi(sigs):
    # for each signature in sigs, get the corresponding phi
    wrapped = sigs.reshape(-1, 2, 3, 16)
    phi = wrapped.sum(2).reshape(-1,32)
    return phi

def get_eta(sigs):
    # for each signature in sigs, get the corresponding eta
    wrapped = sigs.reshape(-1, 6, 16)
    eta = wrapped.sum(2).reshape(-1,3)
    # notrmalize such that etaC and etaT sum to 1 respectively.
    eta = (eta/eta.sum(1)[:,None]).reshape(-1,6)
    return eta

def flatten_eta(eta): # eta pkm -> kc
    warnings.warn('Eta no longer constructed as pkm - use reshape instead', DeprecationWarning)
    return np.moveaxis(eta,0,1).reshape(-1, 6)

def alr(x, e=1e-12):
    '''
    additive log ratio
    x is a NxD matrix
    >>> x = np.array([.1, .3, .4, .2])
    >>> alr(x)
    array([ 1.09861229,  1.38629436,  0.69314718])
    '''
    # add small value for stability in log
    x = x + e
    return (np.log(x) - np.log(x[...,-1]).reshape(-1,1))[:,0:-1]

def alr_inv(y):
    '''
    inverse alr transform
    y is a Nx(D-1) matrix
    >>> x = np.array([.1, .3, .4, .2])
    >>> alr_inv(alr(x))
    array([ 0.1,  0.3,  0.4,  0.2])
    '''
    if y.ndim == 1: y = y.reshape(1,-1)
    return softmax(np.hstack([y, np.zeros((y.shape[0], 1)) ]), axis = 1)

def kmeans_alr(data, nsig, rng=np.random.default_rng()):
    #https://github.com/scikit-learn/scikit-learn/issues/16988#issuecomment-817375063
    km = k_means(alr(data), nsig, random_state=np.random.RandomState(rng.bit_generator))
    return alr_inv(km[0])

def alp_B(data, B):
    return (data * np.log(B)).sum() / data.sum()

def mult_ll(x, p):
    # x and p should both be same dimensions; Sx96
    return loggamma(x.sum(1) + 1) - loggamma(x+1).sum(1) + (x * np.log(p)).sum(1)

def lap_B(data, Bs):
    # Bs should be shape DxSx96 where D is the number of posterior samples
    # use logsumexp for stability
    assert Bs.ndim == 3, 'expected multiple trials for B'
    return logsumexp(np.vstack([mult_ll(data, B) for B in Bs]).sum(1)) - np.log(Bs.shape[0])

def profile_sigs(sigs, refsigs, thresh = 0.9, refidx = None):
    # return mapping of sigs to similar refsigs

    if isinstance(refsigs, pd.core.frame.DataFrame) and refsigs.index is not None:
        refidx = refsigs.index
    elif refidx is None:
        refidx = pd.Index([f'refsig_{i}' for i in range(0,refsigs.shape[0])])

    sim = cosine_similarity(sigs, refsigs)
    closest_dist = np.max(sim, axis=1)
    closest = refidx[np.argmax(sim, axis=1)]
    above_thresh = [str(refidx[x].to_numpy()) for x in sim > thresh]

    df=pd.DataFrame.from_dict({'inferred signature': [f'sig_{i}' for i in range(0,sigs.shape[0])],
                               'closest reference signature':closest,
                               'dist to closest': closest_dist,
                               f'reference signatures with >{thresh} similarity': above_thresh
                              })
    df.index = [f'sig_{i}' for i in range(0,sigs.shape[0])]
    return df

def load_cosmic_V3():
    """Return a dataframe of COSMIC V3 signature definitions

    Contains:
        96 mutation type columns of non-null float64
        78 rows of signature definitions, rows sum to 1

    """
    # This is a stream-like object. If you want the actual info, call
    # stream.read()
    f = pkg_resources.resource_filename(__name__, 'data/COSMIC_v3.2_SBS_GRCh37.txt')
    return load_sigs(f)

def load_default_config():
    """Return a default configuration dict

    Contains:
        96 mutation type columns of non-null float64
        78 rows of signature definitions, rows sum to 1

    """
    # This is a stream-like object. If you want the actual info, call
    # stream.read()
    f = pkg_resources.resource_filename(__name__, 'config/default.yaml')
    return load_config(f)


import errno              # handle errors
import logging            # log errors
import os
import sys
from pathlib import Path  # filesystem related stuff
import re                 # regular expression to find resolution from

import numpy as np
from astropy.wcs import WCS    # handle astronomic coordinates
from astropy.io import fits    # read data from .fits files.
from astropy.io import ascii
from astropy.coordinates import SkyCoord, match_coordinates_sky, Angle
from reproject import reproject_interp

from astropy.visualization import simple_norm
import matplotlib.pyplot as plt

from .constants import two_column, arcsec_to_pixel

logger = logging.getLogger(__name__)

class ParameterClass:
    def __init__(self,name,**kwargs):
        self.name     = name
        for k,v in kwargs.items():
            setattr(self,k,v)

class ReadLineMaps:
    '''load a fits file from the PHANGS--MUSE DAP (the line maps)

    Read the linemaps from the PHANGS-MUSE DAP. You only need to specify
    the folder and the name of the galaxy. The function will pick the
    correct file and read all the specified extensions. If the file is
    from the copt, it will detect the resolution from the filename and
    save it.

    It will also try to read some auxiliary files that should be in
    folder named `AUXILIARY` (next to `folder`). This folder should
    contain star masks, PSF maps (FWHM of each pointing) and an alternative
    [OIII]5007 map (not measured from a fit).
    '''

    def __init__(self,folder,name,extensions=[],**kwargs):
        '''
        Parameters
        ----------

        folder : string
            name of the folder with the data for one object. This folder
            must contain a file with name "FolderName_MAPS.fits" and
            possibly some additional files

        extensions : list
            list of extensions that are read. Each element must be a valid
            extension in the previously defined fits file (the actual name
            of the extension is `ExtensionName_FLUX` and
            `ExtensionName_FLUX_ERR` but they are automaticly completed).

        **kwargs :
            any additional properties of the galaxy (like E(B-V) or
            parameters used during the analysis). They are saved as an
            attribute under the given name.
        '''

        # we simply use the first file in the folder that starts with name
        self.name     = name
        self.filename = next(folder.glob(f'{name}*.fits'))
        self.copt_res = np.float(next(iter(re.findall('-(.*)asec',self.filename.stem)), 'nan'))

        logger.info(f'loading {self.filename.name}')

        self.lines    = []

        # we save the additional parameters
        for k,v in kwargs.items():
            setattr(self,k,v)

        if not self.filename.is_file():
            raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), self.filename)

        # make sure lines is a list
        lines = [extensions] if not isinstance(extensions, list) else extensions

        #==============================================================
        # load the main DAP products
        #==============================================================
        with fits.open(self.filename) as hdul:

            # if no lines are given, we read in all lines
            if len(lines)==0:
                lst = hdul.info(output=False)
                lines = [x[1].split('_')[0] for x in lst if x[1].endswith('_FLUX')]

            # save the white-light image
            header = hdul[f'FLUX'].header
            setattr(self,'header',header)
            setattr(self,'wcs',WCS(header))
            setattr(self,'shape',(header['NAXIS2'],header['NAXIS1']))
            setattr(self,'V_STARS',hdul['V_STARS'].data)
            setattr(self,'stellar_mass',hdul['STELLAR_MASS_DENSITY'].data)
            setattr(self,'stellar_mass_err',hdul['STELLAR_MASS_DENSITY_err'].data)
            setattr(self,'Ebv_stars',hdul['EBV_STARS'].data)
            setattr(self,'whitelight',hdul['FLUX'].data)
            setattr(self,'whitelight_err',hdul['SNR'].data)

            for line in lines:
                # save the main data and the associated error
                setattr(self,line,hdul[f'{line}_FLUX'].data)
                setattr(self,f'{line}_err',hdul[f'{line}_FLUX_ERR'].data)
                setattr(self,f'{line}_SIGMA',np.sqrt(hdul[f'{line}_SIGMA'].data**2 - hdul[f'{line}_SIGMA_CORR'].data**2))
                setattr(self,f'{line}_SIGMA_ERR',hdul[f'{line}_SIGMA_ERR'])

                # append to list of available lines
                self.lines.append(line)

        #==============================================================
        # load auxiliary maps
        #==============================================================

        # and one where OIII is not measured by fitting
        OIII_map_file = folder.parent / 'AUXILIARY' / 'oiii_from_cubes' / f'{self.name}_oiii_flux.fits'
        if OIII_map_file.is_file():
            logger.info(f'replacing OIII5006 map')
            # replace the old line maps with the new one
            setattr(self,'OIII5006_DAP',getattr(self,'OIII5006'))
            setattr(self,'OIII5006_DAP_err',getattr(self,'OIII5006_err'))
            data = fits.getdata(OIII_map_file,0)
            setattr(self,'OIII5006',data)
        else:
            logger.warn(f'"{self.name}_oiii_flux.fits" does not exists.')

        # star mask and seeing map (for PSF)
        star_mask_file = folder.parent / 'AUXILIARY' / 'starmasks' / f'{self.name}_starmask.fits'
        seeing_map_file = folder.parent / 'AUXILIARY' / 'seeing_maps' / f'{self.name}_seeing.fits'

        for filename, description in zip([star_mask_file,seeing_map_file],["star_mask","PSF"]):

            if filename.is_file():
                with fits.open(filename) as hdul:
                    data   = hdul[0].data

                    if self.shape != data.shape:
                        logger.warning(f'{description} map has different shape. Reprojecting')
                        data,_ = reproject_interp(hdul,self.header)

                        # star_mask ist 0 or 1 (even for interpolated pixels)
                        if description=='star_mask':
                            data = np.round(data,0)
                        elif description=='PSF':
                            data = np.round(data,2)

                    setattr(self,description,data)

            else:
                logger.warning(f'no {description} available')

        if not hasattr(self,'star_mask'):
            self.star_mask = np.zeros_like(self.OIII5006)

        if not hasattr(self,'PSF'):
            # for DR2 galaxies where no PSF data exists we assume FWHM=1" for all pointings
            logger.warning('creating 1" seeing map')
            self.PSF = np.ones_like(self.OIII5006)
            self.PSF[np.isnan(self.OIII5006)] = np.nan
        self.PSF *= arcsec_to_pixel
        logger.info(f'galaxy has {len(np.unique(self.PSF[~np.isnan(self.PSF)]))} pointings')

        logger.info(f'file loaded with {len(self.lines)} extensions')

    def __repr__(self):
        '''create an overview of the available attributes'''

        string = ''
        for k,v in self.__dict__.items():
            if type(v) in [str,int,float]:
                string += f'{k}: {v}\n'
            else:
                string += k + '\n'

        return string

    def plot(self,line):
        '''plot a single emission line'''

        if not hasattr(self,line):
            raise AttributeError(f'Object has no map {line}')

        data = getattr(self,line)

        fig = plt.figure(figsize=(two_column,two_column/1.618))
        ax  = fig.add_subplot(projection=self.wcs)

        norm = simple_norm(data,clip=False,percent=99)
        ax.imshow(data,norm=norm)
        plt.show()

        #return fig

# save lines to individual .fits file
def split_fits(filename,extensions):
    '''

    Parameters
    ----------
    filename:
        a fits file containing lines in multiple extensions

    extensions:
        the extensions to save as single files
    '''

    # make sure lines is a list
    extensions = [extensions] if not isinstance(extensions, list) else extensions

    if not os.path.isfile(filename):
        raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), filename)
    print(f'splitting {filename} into seperate files')

    with fits.open(filename) as hdul:
        for ext in extensions:
            data = hdul[f'{ext}']
            data.writeto(f'{ext}.fits',overwrite=True)

    print(f'{len(extensions)} extension(s) saved to seperate file(s)')

class ReadMosaicFiles:

    def __init__(self,filename):
        '''open the large MOSAIC files in python

        the MOSAIC files contain the full spectral information
        '''

        logger.warning('not yet implemented')
        sys.exit()

        with fits.open(filename,memmap=True,mode='denywrite') as hdul:
            wcs = WCS(hdul[1].header)
            data = hdul[1].data

            print(data.shape)
            print(hdul[1].header)
            #data = hdul[f'{line}_FLUX']
            #data.writeto(f'{line}.fits',overwrite=True)

def write_table(table,name,filename):
    '''write the table to a file

    this will create a `.tex` file and a machine readable file

    to read this file use
    f = basedir / 'data' / 'catalogues' / f'{galaxy.name}_{typ}_candidates.txt'
    ascii.read(f,format='fixed_width_two_line',position_char='=')
    '''

    threshold = '0.7"'

    table['RA'],table['DEC'] = zip(*[x.split(' ') for x in table['SkyCoord'].to_string(style='hmsdms',precision=2)])
    table['Galaxy'] = name

    if name =='NGC0628':
        logging.info('comparing to existing studies')
        from .load_references import NGC628
        cat = {'Kreckel PN':'K17;','Herrmann PN':'H08;','Kreckel SNR':'K17s;'}
        ID, angle, Quantity  = match_coordinates_sky(NGC628['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,NGC628):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += cat[row['source']]

    if name =='NGC3351':
        logging.info('comparing to existing studies')
        from .load_references import pn_NGC3351_ciardullo
        ID, angle, Quantity  = match_coordinates_sky(pn_NGC3351_ciardullo['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,pn_NGC3351_ciardullo):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += 'C02;'

    if name =='NGC3627':
        logging.info('comparing to existing studies')
        from .load_references import pn_NGC3627_ciardullo
        ID, angle, Quantity  = match_coordinates_sky(pn_NGC3627_ciardullo['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,pn_NGC3627_ciardullo):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += 'C02;'

    if name =='NGC5068':
        logging.info('comparing to existing studies')
        from .load_references import pn_NGC5068_herrmann
        ID, angle, Quantity  = match_coordinates_sky(pn_NGC5068_herrmann['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,pn_NGC5068_herrmann):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += 'H08;'

    for col in ['OIII5006','HA6562','NII6583','SII']:
        table[col][np.where(table[col]<table[f'{col}_err'])] = table[f'{col}_err'][np.where(table[col]<table[f'{col}_err'])]

    table['OIII/Ha']   = table['OIII5006']/table['HA6562']
    table['d(OIII/Ha)']= table['OIII5006'] / table['HA6562'] * np.sqrt( (table['HA6562_err'] / table['HA6562'])**2 + (table['OIII5006_err'] / table['OIII5006'])**2)

    table['Ha/NII']    = table['HA6562']/table['NII6583']
    table['d(Ha/NII)'] = table['HA6562'] / table['NII6583'] * np.sqrt( (table['NII6583_err'] / table['NII6583'])**2 + (table['HA6562_err'] / table['HA6562'])**2)

    table['Ha/SII']    = table['HA6562']/table['SII']
    table['d(Ha/SII)'] = table['HA6562'] / table['SII'] * np.sqrt( (table['SII_err'] / table['SII'])**2 + (table['HA6562_err'] / table['HA6562'])**2)

    for col in ['mOIII','dmOIII','v_SIGMA','OIII/Ha','d(OIII/Ha)','Ha/NII','d(Ha/NII)','Ha/SII','d(Ha/SII)']:
        table[col].info.format = '%.2f'

    #
    for typ in ['PN','SNR']:

        tbl_out = table[table['type']==typ]

        if typ == 'PN':
            n = 'Planetary Nebula'
        if typ == 'SNR':
            n = 'Supernova Remnants'

        tbl_out.sort('mOIII')
        tbl_out['name'] = np.arange(1,len(tbl_out)+1)

        # add marker for existing study or excluded object
        notes = []
        for i,row in enumerate(tbl_out):
            note = ''
            if name == 'NGC0628' or name=='NGC3351' or name=='NGC3627' or name=='NGC5068':
                note += row['match']
            if row['overluminous']:
                note += 'OL;'
            if row['SNRorPN']:
                note += 'PN;'
            notes.append(note)
        tbl_out['notes'] = [x.strip(';') for x in notes]

        tbl_out.rename_columns(['name','RA','DEC','v_SIGMA'],['ID','R.A.','Dec.','sigmaV'])
        tbl_out = tbl_out[['Galaxy','ID','notes','R.A.','Dec.','mOIII','dmOIII','OIII/Ha','d(OIII/Ha)',
                           'Ha/NII','d(Ha/NII)','Ha/SII','d(Ha/SII)','sigmaV']]

        with open((filename / f'{name}_{typ}_candidates').with_suffix('.txt'),'w',newline='\n') as f:
            ascii.write(tbl_out,f,format='fixed_width_two_line',overwrite=True,delimiter_pad=' ',position_char='=')

        with open((filename / f'{name}_{typ}_candidates').with_suffix('.tex'),'w',newline='\n') as f:
            ascii.write(tbl_out,f,Writer=ascii.Latex,overwrite=True)

    logger.info(f'table saved to files (for {name})')

def write_LaTeX_old(table,galaxy,filename):
    '''write the table to a file

    this will create a `.tex` file and a machine readable file

    to read this file use
    f = basedir / 'data' / 'catalogues' / f'{galaxy.name}_{typ}_candidates.txt'
    ascii.read(f,format='fixed_width_two_line',position_char='=')
    '''

    threshold = '0.7"'

    table['SkyCoord'] = SkyCoord.from_pixel(table['x'],table['y'],galaxy.wcs)
    table['RA'],table['DEC'] = zip(*[x.split(' ') for x in table['SkyCoord'].to_string(style='hmsdms',precision=2)])

    if galaxy.name =='NGC0628':
        print('comparing to existing studies')
        from .load_references import NGC628
        cat = {'Kreckel PN':'K17','Herrmann PN':'H08','Kreckel SNR':'K17s'}
        ID, angle, Quantity  = match_coordinates_sky(NGC628['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,NGC628):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += cat[row['source']]

    if galaxy.name =='NGC3351':
        print('comparing to existing studies')
        from .load_references import pn_NGC3351_ciardullo
        ID, angle, Quantity  = match_coordinates_sky(pn_NGC3351_ciardullo['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,pn_NGC3351_ciardullo):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += 'C02'

    if galaxy.name =='NGC3627':
        print('comparing to existing studies')
        from .load_references import pn_NGC3627_ciardullo
        ID, angle, Quantity  = match_coordinates_sky(pn_NGC3627_ciardullo['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,pn_NGC3627_ciardullo):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += 'C02'

    if galaxy.name =='NGC5068':
        print('comparing to existing studies')
        from .load_references import pn_NGC5068_herrmann
        ID, angle, Quantity  = match_coordinates_sky(pn_NGC5068_herrmann['SkyCoord'],table['SkyCoord'])
        table['match'] = np.empty(len(table),dtype='U12')
        for i,a,row in zip(ID,angle,pn_NGC5068_herrmann):
            if a.__lt__(Angle('0.8"')):
                table['match'][i] += 'H08'

    for typ in ['PN','SNR']:

        if typ == 'PN':
            n = 'Planetary Nebula'
        if typ == 'SNR':
            n = 'Supernova Remnants'
        if typ == 'HII':
            n = '\\HII-regions'

        latexdict = {'tabletype': 'table*',
        'header_start': '\\toprule\\toprule',
        'header_end': '\\midrule',
        'data_end': '\\bottomrule',
        'caption': f'{n} Identifications',
        'units': {'R.A.':'(J2000)','Dec.':'(J2000)','$m_\\OIII$':'mag','d$m_\\OIII$':'mag','$\sigma_V$':'\\si{\\km \per \\second}'},
        'preamble': '\\centering',
        'tablefoot': f'\\label{{tbl:{typ}_Identifications}}'
                    }

        tbl_out = table[table['type']==typ]
        tbl_out['Galaxy'] = galaxy.name
        tbl_out.sort('mOIII')

        tbl_out['OIII/Ha']   = np.empty(len(tbl_out),dtype='U8')
        tbl_out['Ha/NII']    = np.empty(len(tbl_out),dtype='U8')
        tbl_out['Ha/SII']    = np.empty(len(tbl_out),dtype='U8')
        tbl_out['d(OIII/Ha)']= np.empty(len(tbl_out),dtype='U8')
        tbl_out['d(Ha/NII)'] = np.empty(len(tbl_out),dtype='U8')
        tbl_out['d(Ha/SII)'] = np.empty(len(tbl_out),dtype='U8')

        # add marker for existing study or excluded object
        names = np.arange(1,len(tbl_out)+1)
        notes = []
        for i,row in enumerate(tbl_out):
            notes = ''
            if galaxy.name == 'NGC0628' or galaxy.name=='NGC3351' or galaxy.name=='NGC3627' or galaxy.name=='NGC5068':
                name += row['match']
            if row['overluminous']:
                name += '+'
            if row['SNRorPN']:
                name += 'PN'

            names.append(name)
        tbl_out['name'] = names
        tbl_out['notes'] = notes

        # calculate line ratios with limits (> sign)
        for i,row in enumerate(tbl_out):
            if not row['HA6562_detection']:
                row['OIII/Ha'] += '>'
            row['OIII/Ha'] += f"{row['OIII5006'] / row['HA6562']:.2f}"
            row['d(OIII/Ha)'] = f"{row['OIII5006'] / row['HA6562'] * np.sqrt( (row['HA6562_err'] / row['HA6562'])**2 + (row['OIII5006_err'] / row['OIII5006'])**2):.2f}"

            if not row['HA6562_detection'] and not row['NII6583_detection']:
                row['Ha/NII'] = '...'
                row['d(Ha/NII)'] = '...'
            elif not row['HA6562_detection']:
                row['Ha/NII'] = '<'
                row['Ha/NII'] += f"{row['HA6562'] / row['NII6583']:.2f}"
                row['d(Ha/NII)'] = f"{row['HA6562'] / row['NII6583'] * np.sqrt( (row['NII6583_err'] / row['NII6583'])**2 + (row['HA6562_err'] / row['HA6562'])**2):.2f}"
            elif not row['NII6583_detection']:
                row['Ha/NII'] = '>'
                row['Ha/NII'] += f"{row['HA6562'] / row['NII6583']:.2f}"
                row['d(Ha/NII)'] = f"{row['HA6562'] / row['NII6583'] * np.sqrt( (row['NII6583_err'] / row['NII6583'])**2 + (row['HA6562_err'] / row['HA6562'])**2):.2f}"
            else:
                row['Ha/NII'] += f"{row['HA6562'] / row['NII6583']:.2f}"
                row['d(Ha/NII)'] = f"{row['HA6562'] / row['NII6583'] * np.sqrt( (row['NII6583_err'] / row['NII6583'])**2 + (row['HA6562_err'] / row['HA6562'])**2):.2f}"

            if not row['HA6562_detection'] and not row['SII_detection']:
                row['Ha/SII'] = '...'
                row['d(Ha/SII)'] = '...'
            elif not row['HA6562_detection']:
                row['Ha/SII'] += '<'
                row['Ha/SII'] += f"{row['HA6562'] / row['SII']:.2f}"
                row['d(Ha/SII)'] = f"{row['HA6562'] / row['SII'] * np.sqrt( (row['HA6562_err'] / row['HA6562'])**2 + (row['SII_err'] / row['SII'])**2):.2f}"
            elif not row['SII_detection']:
                row['Ha/SII'] += '>'
                row['Ha/SII'] += f"{row['HA6562'] / row['SII']:.2f}"
                row['d(Ha/SII)'] = f"{row['HA6562'] / row['SII'] * np.sqrt( (row['HA6562_err'] / row['HA6562'])**2 + (row['SII_err'] / row['SII'])**2):.2f}"
            else:
                row['Ha/SII'] += f"{row['HA6562'] / row['SII']:.2f}"
                row['d(Ha/SII)'] = f"{row['HA6562'] / row['SII'] * np.sqrt( (row['HA6562_err'] / row['HA6562'])**2 + (row['SII_err'] / row['SII'])**2):.2f}"

        tbl_out['mOIII'].info.format = '%.2f'
        tbl_out['dmOIII'].info.format = '%.2f'
        tbl_out['v_SIGMA'].info.format = '%.2f'

        tbl_out.rename_columns(['name','RA','DEC','v_SIGMA'],['ID','R.A.','Dec.','sigmaV'])
        tbl_out = tbl_out[['Galaxy','ID','R.A.','Dec.','mOIII','dmOIII','OIII/Ha','d(OIII/Ha)',
                           'Ha/NII','d(Ha/NII)','Ha/SII','d(Ha/SII)','sigmaV']]

        with open((filename / f'{galaxy.name}_{typ}_candidates').with_suffix('.txt'),'w',newline='\n') as f:
            ascii.write(tbl_out,f,format='fixed_width_two_line',overwrite=True,delimiter_pad=' ',position_char='=')

        mOIII  = []
        OIIIHA = []
        HANII  = []
        HASII  = []
        for row in tbl_out:
            if row["mOIII"] == '...':
                mOIII.append('...')
            else:
                mOIII.append(f'{row["mOIII"]:.2f} $\pm$ {row["dmOIII"]:.2f}')
            if row["OIII/Ha"] == '...':
                OIIIHA.append('...')
            else:
                OIIIHA.append(f'{row["OIII/Ha"]} $\pm$ {row["d(OIII/Ha)"]}')
            if row["Ha/NII"] == '...':
                HANII.append('...')
            else:
               HANII.append(f'{row["Ha/NII"]} $\pm$ {row["d(Ha/NII)"]}')
            if row["Ha/SII"] == '...':
                HASII.append('...')
            else:
               HASII.append(f'{row["Ha/SII"]} $\pm$ {row["d(Ha/SII)"]}')

        tbl_out['$m_\\OIII$'] = mOIII
        tbl_out['$\\OIII/\\HA$'] = OIIIHA
        tbl_out['$\\HA/\\NII$'] = HANII
        tbl_out['$\\HA/\\SII$'] = HASII
        tbl_out.rename_column('sigmaV',f'$\sigma_V$')
        tbl_out = tbl_out[['ID','R.A.','Dec.','$m_\\OIII$','$\\OIII/\\HA$',
                           '$\\HA/\\NII$','$\\HA/\\SII$',f'$\sigma_V$']]

        with open((filename / f'{galaxy.name}_{typ}_candidates').with_suffix('.tex'),'w',newline='\n') as f:
            ascii.write(tbl_out,f,Writer=ascii.Latex, latexdict=latexdict,overwrite=True)

        # shorten column names for machine readable table
        #for col in tbl_out.colnames:
        #    tbl_out.rename_column(col,col.translate({ord(s): None for s in '$\_'}))

    logger.info(f'table saved to files (for {galaxy.name})')

def read_catalogue(filename):

    catalogue = ascii.read(filename,format='fixed_width_two_line',delimiter_pad=' ',position_char='=')
    catalogue['SkyCoord'] = SkyCoord(catalogue['R.A.'],catalogue['Dec.'])

    return catalogue


# -*- coding: utf-8 -*-
# @Author  : Daniel Ordonez
# @email   : daniels.ordonez@gmail.com

from typing import Optional

import numpy as np
from numpy import ndarray

from .slip_model import SlipModel, THETA, X

class SlipGaitCycle(object):

    def __init__(self, slip_model: SlipModel, t_flight: ndarray, flight_cartesian_traj: ndarray, t_stance: ndarray,
                 stance_polar_traj: ndarray, target_to_state: Optional[ndarray] = None,
                 optimization_cost: Optional[float] = np.NaN):
        """
        Class representing a Slip gait cycle composed of:
         - a flight phase: starting with the take-off (TO) event of a previous gait cycle and ending at the present cycle
          touch-down (TD) event. And,
         - a stance phase: starting at the present cycle TD event and ending with the cycle's TO event.
        It is assumed the gait cycle starts with a flight phase and ends with the stance phase.
        :param slip_model: SLIP model providing the `m`, `k` and `r0` parameters.
        :param t_flight: (F,) Discrete time array during flight phase
        :param flight_cartesian_traj: (6, F) Cartesian state [x, x', x'', z, z', z''] of the SLIP CoM at each time
        during the flight phase
        :param t_stance: (S,) Discrete time array during flight phase
        :param stance_polar_traj: (4, S) Polar state [theta,theta',r,r'] of the SLIP CoM at each time during
        the stance phase, in a reference frame centered at the foot contact point with the ground.
        :param target_to_state: (6,) Optional cartesian target take-off state, used for control. Use np.NaN in the
        dimensions of the cartesian state that are irrelevant for control (e.g. z''=np.NaN)
        :param optimization_cost: Optional scalar indicating the optimization cost of the gait cycle. By default is set
        to be the euclidean norm of the error between the real TO state and the target TO state (ignoring np.NaNs dims)
        """
        self.slip_model = slip_model
        if not flight_cartesian_traj.shape[0] == 6:
            raise AssertionError("Expected cartesian trajectory [x, x', x'', z, z', z'']")
        else:
            if not len(t_flight) == flight_cartesian_traj.shape[1]:
                raise AssertionError('Invalid flight trajectory')
            else:
                self.t_flight = t_flight
                self.flight_cartesian_traj = flight_cartesian_traj
                assert stance_polar_traj.shape[0] == 4, "Expected stance polar trajectory [theta, theta', r, r']"
                assert len(t_stance) == stance_polar_traj.shape[1], 'Invalid flight trajectory'
            self.t_stance = t_stance
            self.stance_polar_traj = stance_polar_traj
            self._stance_cartesian_traj = None
            self.touch_down_angle = self.stance_polar_traj[(THETA, 0)]
            self.take_off_angle = self.stance_polar_traj[(THETA, -1)]
            self.foot_contact_pos = self.flight_cartesian_traj[(X, -1)] - self.slip_model.r0 * np.cos(
                self.touch_down_angle)
            assert self.t_flight[(-1)] == self.t_stance[
                0], 'Touch down state should be in flight and stance trajectories'
        self.target_to_state = np.array(target_to_state) if target_to_state is not None else np.ones((6,)) * np.NaN
        self.optimization_cost = optimization_cost
        if not np.all(np.isnan(self.target_to_state)):
            self.target_to_state = np.ma.array((self.target_to_state), mask=(np.isnan(self.target_to_state)))
            if np.isnan(optimization_cost):
                self.optimization_cost = np.linalg.norm(self.take_off_state - target_to_state)

    @property
    def stance_cartesian_traj(self) -> ndarray:
        if self._stance_cartesian_traj is None:
            self._stance_cartesian_traj = self.slip_model.polar_to_cartesian(trajectory=(self.stance_polar_traj),
                                                                             foot_contact_pos=(self.foot_contact_pos))
        return self._stance_cartesian_traj

    @property
    def take_off_state(self) -> ndarray:
        if self._stance_cartesian_traj is None:
            return self.slip_model.polar_to_cartesian((self.stance_polar_traj[:, -1]),
                                                      foot_contact_pos=(self.foot_contact_pos))
        else:
            return np.array(self.stance_cartesian_traj[:, -1])

    @property
    def take_off_state_polar(self):
        return np.array(self.stance_polar_traj[:, -1])

    @property
    def touch_down_state(self):
        if self._stance_cartesian_traj is None:
            return self.slip_model.polar_to_cartesian((self.stance_polar_traj[:, 0]),
                                                      foot_contact_pos=(self.foot_contact_pos))
        else:
            return np.array(self.stance_cartesian_traj[:, 0])

    @property
    def touch_down_state_polar(self):
        return np.array(self.stance_polar_traj[:, 0])

    @property
    def prev_take_off_state(self):
        return np.array(self.flight_cartesian_traj[:, 0])

    @property
    def start_time(self):
        return float(self.t_flight[0])

    @property
    def end_time(self):
        return float(self.t_stance[(-1)])

    def offset_initial_time(self, time_offset):
        self.t_flight += time_offset
        self.t_stance += time_offset

    def __str__(self):
        return 'Cost:%.2f TD:%.1f[deg] TO:%.1f[deg] time:[%.2f, %.2f]' % (
            self.optimization_cost, np.rad2deg(self.touch_down_angle), np.rad2deg(self.take_off_angle),
            self.start_time, self.end_time)

    def __repr__(self):
        return str(self)

class SlipGaitCycleCtrl(SlipGaitCycle):

    def __init__(self, slip_model, t_flight, flight_cartesian_traj, t_stance, stance_passive_polar_traj=None,
                 stance_ctrl_polar_traj=None, control_signal=None, target_to_state=None, optimization_cost=np.NaN,
                 ctrl_kwargs=None):
        """
        Class representing a Controlled Slip gait Cycle. This assumes the SLIP model is an actuated extended version
        (see "Learning to run naturally: Guiding policies with the Spring-Loaded Inverted Pendulum" Chap 4.1) where the
        control inputs are a resting leg length displacement `r_delta` (axial force control) and a hip torque `tau_hip`.
        The main difference with the parent class it that `SlipGaitCycleCtrl` stores also a `control_signal` and if
        provided an additional `SlipGaitCycle` instance representing the passive dynamical response of SLIP (useful for
        plotting animation and intuition).
        :param slip_model: SLIP model providing the `m`, `k` and `r0` parameters.
        :param t_flight: (F,) Discrete time array during flight phase
        :param flight_cartesian_traj: (6, F) Cartesian state [x,x',x'',z,z',z''] of the SLIP CoM at each time during
        the flight phase
        :param t_stance: (S,) Discrete time array during flight phase
        :param stance_passive_polar_traj: (4, S) Passive Polar state [theta,theta',r,r'] of the SLIP CoM at each time
        during the stance phase, in a reference frame centered at the foot contact point with the ground.
        :param stance_ctrl_polar_traj: (4, S) Controlled Polar state [theta,theta',r,r'] of the SLIP CoM at each time
        during the stance phase, in a reference frame centered at the foot contact point with the ground.
        :param control_signal: (2, S) Control input signal assumed to hold ()
        :param target_to_state: (6,) Optional cartesian target take-off state, used for control. Use np.NaN in the
        dimensions of the cartesian state that are irrelevant for control (e.g. z''=np.NaN)
        :param optimization_cost: Optional scalar indicating the optimization cost of the gait cycle. By default is set
        to be the euclidean norm of the error between the real TO state and the target TO state (ignoring np.NaNs dims)
        :param ctrl_kwargs: Dictionary holding controller-related keyword arguments.
        """
        self.control_signal = control_signal
        self.optimization_cost = optimization_cost
        self.control_signal = control_signal
        self.ctrl_kwargs = ctrl_kwargs
        super(SlipGaitCycleCtrl, self).__init__(slip_model, t_flight, flight_cartesian_traj, t_stance,
                                                stance_polar_traj=stance_ctrl_polar_traj,
                                                target_to_state=target_to_state)
        if stance_passive_polar_traj is not None:
            self.passive_gait_cycle = SlipGaitCycle(slip_model, t_flight, flight_cartesian_traj, t_stance,
                                                    stance_polar_traj=stance_passive_polar_traj)
        else:
            self.passive_gait_cycle = None

    @property
    def stance_cartesian_traj(self) -> ndarray:
        if self._stance_cartesian_traj is None:
            self._stance_cartesian_traj = self.slip_model.polar_to_cartesian(trajectory=self.stance_polar_traj,
                                                                             control_signal=self.control_signal,
                                                                             foot_contact_pos=self.foot_contact_pos)
        return self._stance_cartesian_traj

    @property
    def take_off_state(self) -> ndarray:
        if self._stance_cartesian_traj is None:
            return self.slip_model.polar_to_cartesian(self.stance_polar_traj[:, -1],
                                                      control_signal=self.control_signal[:, -1],
                                                      foot_contact_pos=self.foot_contact_pos)
        else:
            return np.array(self.stance_cartesian_traj[:, -1])

    @property
    def touch_down_state(self):
        if self._stance_cartesian_traj is None:
            return self.slip_model.polar_to_cartesian(self.stance_polar_traj[:, 0],
                                                      control_signal=self.control_signal[:, 0],
                                                      foot_contact_pos=self.foot_contact_pos)
        else:
            return np.array(self.stance_cartesian_traj[:, 0])


# -*- coding: utf-8 -*-
import theano
import theano.tensor as T
import numpy as np
import pandas as pd
from collections import OrderedDict

INIT_RANGE = 0.01

def uniform(lb, ub, size):
    """

    :param lb: lower bound
    :param ub: upper bound
    :param size: tensor shape
    :return:
    """
    return np.array(np.random.uniform(low=lb, high=ub, size=size), dtype='float32')

def glorot_uniform(size):
    """
    glorot uniform initializer
    """
    if len(size) == 1:
        values = np.zeros_like(size)
    else:
        scale = np.sqrt(6.0 / np.sum(size))
        values = np.random.uniform(low=-scale, high=scale, size=size)
    return values.astype('float32')

def orthogonal(size):
    """
    equivalent to orthogonal_init but return numpy array
    """
    if len(size) == 1:
        values = np.zeros_like(size)
    else:
        a = np.random.normal(loc=0.0, scale=1.0, size=size)
        # reconstruction based on reduced SVD
        u, _, v = np.linalg.svd(a, full_matrices=False)
        q = u if u.shape == size else v
        q = q.reshape(size)
        values = q
    return values.astype("float32")

def zeros(size):
    """
    generate zero vector / tensor
    :param size: size
    :return:
    """
    return np.array(np.zeros(shape=size), dtype='float32')

def ones(size):
    """
    generate one vector / tensor
    :param size:
    :return:
    """
    return np.array(np.zeros(shape=size), dtype='float32')

def lstm_init(n_in, n_out, component="LSTM"):
    """

    :param n_in: input size
    :param n_out: hidden size
    :param component: component name
    :return:
    """
    if True:
        print("Initialize LSTM weights from uniform distribution...")
        W_values = np.concatenate([uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_in, n_out)),
                                   uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_in, n_out)),
                                   uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_in, n_out)),
                                   uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_in, n_out))], axis=1)

        U_values = np.concatenate([uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_out, n_out)),
                                   uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_out, n_out)),
                                   uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_out, n_out)),
                                   uniform(lb=-INIT_RANGE, ub=INIT_RANGE, size=(n_out, n_out))], axis=1)
        b_values = np.concatenate([zeros(n_out),
                                   zeros(n_out),
                                   zeros(n_out),
                                   zeros(n_out)])
    else:
        print("Initialize LSTM weights from glorot uniform + glorot uniform...")
        W_values = np.concatenate([glorot_uniform(size=(n_in, n_out)),
                                  glorot_uniform(size=(n_in, n_out)),
                                  glorot_uniform(size=(n_in, n_out)),
                                  glorot_uniform(size=(n_in, n_out))], axis=1)

        U_values = np.concatenate([glorot_uniform(size=(n_out, n_out)),
                                   glorot_uniform(size=(n_out, n_out)),
                                   glorot_uniform(size=(n_out, n_out)),
                                   glorot_uniform(size=(n_out, n_out))], axis=1)
        b_values = np.concatenate([zeros(n_out),
                                   ones(n_out),  # set forget gate to 1.0
                                   zeros(n_out),
                                   zeros(n_out)])
    W = theano.shared(value=W_values, name='%s_W' % component)
    U = theano.shared(value=U_values, name='%s_U' % component)
    b = theano.shared(value=b_values, name='%s_b' % component)
    return W, U, b

def reverse_tensor(tensor):
    """

    :param tensor: input tensor, shape: (bs, seq_len, n_in)
    :return:
    """
    new_tensor = tensor.dimshuffle(1, 0, 2)
    return new_tensor[::-1].dimshuffle(1, 0, 2)

def adam(cost, params, max_norm=3.0, lr=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
    """
    adam optimizer, default learning rate is 0.001
    """

    grads = T.grad(cost, params)
    t_prev = theano.shared(value=np.float32(0.0))
    updates = OrderedDict()

    t = t_prev + 1
    one = T.constant(1, dtype='float32')

    a_t = lr * T.sqrt(one - beta2 ** t) / (one - beta1 ** t)

    for p, g in zip(params, grads):
        value = p.get_value(borrow=True)
        m_prev = theano.shared(np.zeros(value.shape, dtype=value.dtype),
                               broadcastable=p.broadcastable)
        v_prev = theano.shared(np.zeros(value.shape, dtype=value.dtype),
                               broadcastable=p.broadcastable)

        m_t = beta1 * m_prev + (one - beta1) * g
        v_t = beta2 * v_prev + (one - beta2) * g ** 2

        step = a_t * m_t / (T.sqrt(v_t) + epsilon)

        stepped_p = p - step
        updates[m_prev] = m_t
        updates[v_prev] = v_t
        # add max norm constraint on the W_hidden
        # if p.name == 'W_hidden':
        #	col_norms = T.sqrt(T.sum(T.sqr(stepped_p), axis=0))
        #	desired_norms = T.clip(col_norms, 0, max_norm)
        #	scale = desired_norms / (1e-7 + col_norms)
        #	updates[p] = stepped_p * scale
        # else:
        #	updates[p] = stepped_p
        updates[p] = stepped_p
    updates[t_prev] = t
    return updates

def sgd_momentum(cost, params, lr=0.01, momentum=0.9):
    """
    sgd with momentum
    :param cost: training loss
    :param params: parameters
    :param lr: learning rate
    :param momentum: amount of momentum to apply
    :return:
    """
    grads = T.grad(cost, params)
    updates = OrderedDict()
    for p, g in zip(params, grads):
        updates[p] = p - lr * g
    for p in params:
        value = p.get_value(borrow=True)
        velocity = theano.shared(np.zeros(value.shape, dtype=value.dtype), broadcastable=p.broadcastable)
        x = momentum * velocity + updates[p]
        updates[velocity] = x - p
        updates[p] = x
    return updates

def get_batch_input(max_sequence_length, dataset, bs, idx):
    """

    :param dataset: dataset
    :param bs: batch size
    :param idx: batch index
    :return:
    """
    batch_input = dataset[idx*bs:(idx+1)*bs]
    #print(batch_input)
    batch_data = pd.DataFrame.from_dict(batch_input)
    #print(batch_data)
    #print(batch_data)
    target_fields = ['wids', 'tids', 'y', 'pw']
    batch_input_var = []
    for key in target_fields:
        #print("for %s data" % key)
        data = list(batch_data[key].values)
        #print(data[0])
        #print(data[5])
        if key == 'pw':
            # for position weights
            #batch_input_var.append(theano.shared(value=np.array(data, dtype='float32')))
            batch_input_var.append(np.array(data, dtype='float32'))
        else:
            #batch_input_var.append(theano.shared(value=np.array(data, dtype='int32')))
            batch_input_var.append(np.array(data, dtype='int32'))
    batch_input_var.append(list(batch_data['sent'].values))

    graph_data = list(batch_data['graphsage_embedding'].values)
    graph_data = np.array(graph_data, dtype='float32')
    graphsage_embedding = []

    for s in range(len(batch_input_var[0])):
        graph_embedding = []
        for i in range(max_sequence_length):
            graph_embedding.append(graph_data[s])
        graphsage_embedding.append(graph_embedding)

    graphsage_embedding = np.array(graphsage_embedding, dtype=np.float32)
    batch_input_var.append(graphsage_embedding)

    return batch_input_var


#!/usr/bin/env python
"""Pre-calculate a grid of 5-sigma limiting mags for all times, filters, and DDF fields
"""

# imports
import sys
import warnings
import logging
from argparse import ArgumentParser

import yaml
import numpy as np
import pandas as pd
from astropy.coordinates import SkyCoord
import astropy.coordinates
from astropy.time import Time

import lsst.sims.utils
from lsst.sims.downtimeModel import ScheduledDowntimeData
from lsst.sims.skybrightness_pre import SkyModelPre
from lsst.sims.seeingModel import SeeingModel

# constants

DEFAULT_CONFIG = {
    "start_time": "2022-11-01T16:00:00Z",
    "end_time": "2033-11-01T16:00:00Z",
    "time_freq": "10min",
    "site_name": "LSST",
    "fields": {
        "Elias S1": SkyCoord("00h37m48s", "-44d00m00s"),
        "XMM-LSS": SkyCoord("02h22m50s", "-04d45m00s"),
        "ECDFS": SkyCoord("03h32m30s", "-28d06m00s"),
        "COSMOS": SkyCoord("10h00m24s", "+02d10m55s"),
        "Euclid 1": SkyCoord("03h55m52.8s", "-49d16m48s"),
        "Euclid 2": SkyCoord("04h14m24s", "-47d36m00s"),
    },
    "max_sun_alt_deg": -18.0,
    "max_field_airmass": 2.6,
}

SKY_MODEL = SkyModelPre()

# exception classes

# interface functions

def calc_m5(config=None):
    """Calculate the 5-sigma limiting magnitudes for a set of fields

    Parameters
    ----------
    config : `dict'
        Configuration parameters, with the following contents:
        start_time : `str` or period-like, default None
            Start of time range
        end_time : `str` or period-like, default None
            End of time range
        time_freq : `str`
            Time frequency, designated by a pandas frequency string
        fields : `dict` [`str`, ~astropy.coordinates.SkyCoord`]
            The pointings for which to calculate limiting mags.
        site_name : `str`
            The name of the observatory site.
        max_sun_alt_deg: `float`
            Maximum altitude of sun for observing time, in degrees
        max_field_airmass: `float`
            Maximum airmass at which to observe

    Returns
    -------
    m5limits : `pandas.DataFrame`
        A `pandas.Series` of the 5 sigma limiting magnitudes, indexed by the `mjd` and field name.

    Notes
    -----

    For more on pandas frequency strings, see `the pandas documentation
<https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases>`__.

    """

    config = DEFAULT_CONFIG if config is None else config

    # Work around bug in lsst.sims.utils.Site, which compares by reference
    # rather than equivilence
    site_name = "LSST" if config["site_name"] == "LSST" else config["site_name"]
    site = lsst.sims.utils.Site(site_name)
    location = astropy.coordinates.EarthLocation(
        lat=site.latitude, lon=site.longitude, height=site.height
    )

    periods = _init_periods(config, location)

    field_coords = _init_field_coords(config)

    logger.debug("Building the sampled_fields DataFrame")
    # Cross join using the hacky pandas method
    periods["dummy"] = True
    field_coords["dummy"] = True
    sampled_fields = (
        periods.merge(field_coords, on="dummy", how="outer")
        .drop(columns=["dummy"])
        .set_index(["field_name", "period"], drop=False)
    )
    sampled_fields.query("observing", inplace=True)

    logger.debug("Calculating field airmass and filtering")
    field_zd_rad = np.pi / 2 - (
        SkyCoord(
            ra=sampled_fields.field_ra,
            dec=sampled_fields.field_decl,
            frame="icrs",
            unit="deg",
        )
        .transform_to(
            astropy.coordinates.AltAz(
                location=location, obstime=Time(sampled_fields["mjd"], format="mjd")
            )
        )
        .alt.rad
    )
    sampled_fields["field_airmass"] = 1.0 / np.cos(field_zd_rad)
    max_field_airmass = config["max_field_airmass"]
    sampled_fields.query(f"field_airmass < {max_field_airmass}", inplace=True)

    logger.debug("Calculating field angle with the moon")
    sampled_fields["moon_angle"] = (
        SkyCoord(
            ra=sampled_fields["field_ra"],
            dec=sampled_fields["field_decl"],
            frame="icrs",
            unit="deg",
        ).separation(
            SkyCoord(
                ra=sampled_fields["moon_ra"],
                dec=sampled_fields["moon_decl"],
                frame="icrs",
                unit="deg",
            )
        )
    ).deg

    logger.debug("Calculating sky brightness")
    sampled_fields = sampled_fields.groupby("mjd").apply(_get_sky_mags)

    sampled_fields = sampled_fields.groupby("band").apply(_compute_band_fwhms)

    sampled_fields = sampled_fields.groupby("band").apply(_compute_band_m5)

    sampled_fields.set_index(["field_name", "period"], drop=False, inplace=True)

    logger.debug("Finished constructing the m5 DataFrame")
    return sampled_fields

# classes

# internal functions & classes

def _init_periods(config, location):

    logger.debug("Laying down time period boundries")
    period_index = pd.period_range(
        config["start_time"],
        config["end_time"],
        freq=config["time_freq"],
        name="period",
    )
    times = Time(
        pd.to_datetime(period_index.to_timestamp()),  # pylint: disable=no-member
        location=location,
    )

    logger.debug("Calculating solar ephemeris")
    sun_coords = astropy.coordinates.get_sun(times)

    logger.debug("Calculating lunar ephemeris")
    moon_coords = astropy.coordinates.get_moon(times)
    moon_elongation = moon_coords.separation(sun_coords)

    logger.debug("Converting solar and lunar coordinates to horizon system")
    horizon_coordinate_system = astropy.coordinates.AltAz(location=location)
    sun_hzn = sun_coords.transform_to(horizon_coordinate_system)
    moon_hzn = moon_coords.transform_to(horizon_coordinate_system)

    logger.debug("Building the period DataFrame")
    periods = pd.DataFrame(
        {
            "mjd": times.mjd,
            "time": times,
            "lst": times.sidereal_time("mean"),
            "sun_ra": sun_coords.ra.deg,
            "sun_decl": sun_coords.dec.deg,
            "sun_alt": sun_hzn.alt.deg,
            "moon_ra": moon_coords.ra.deg,
            "moon_decl": moon_coords.dec.deg,
            "moon_alt": moon_hzn.alt.deg,
            "moon_elongation": moon_elongation.deg,
        },
        index=period_index,
    )

    periods["moon_waxing"] = (
        periods["moon_elongation"].shift(-1) > periods["moon_elongation"]
    )

    previous_waxing = periods["moon_waxing"].shift(
        1, fill_value=periods.iloc[0]["moon_waxing"]
    )
    periods["new_moon"] = periods["moon_waxing"] & ~previous_waxing
    periods["lunation"] = periods["new_moon"].cumsum()

    mean_local_solar_jd = 2400000.5 + periods["mjd"] + (location.lon.deg / 360.0)
    night_local_solar_jd = np.floor(mean_local_solar_jd).astype(int)
    periods["night"] = night_local_solar_jd - np.min(night_local_solar_jd) + 1

    # Mark daytime periods
    periods["observing"] = periods.sun_alt <= config["max_sun_alt_deg"]

    # Mark scheduled downtime
    periods.reset_index(inplace=True)
    periods.set_index("mjd", inplace=True)
    for down_time in ScheduledDowntimeData(times.min())():
        periods.loc[down_time["start"].mjd : down_time["end"].mjd, "observing"] = False
    periods.reset_index(inplace=True)
    periods.set_index("period", drop=False, inplace=True)
    return periods

def _init_field_coords(config):
    fields = config["fields"]

    logger.debug("Building the fields DataFrame")
    field_coords = (
        pd.DataFrame(
            {
                fld: {
                    "field_name": fld,
                    "field_ra": fields[fld].ra.deg,
                    "field_decl": fields[fld].dec.deg,
                }
                for fld in fields
            }
        )
        .T.set_index("field_name")
        .apply(pd.to_numeric)
        .reset_index()
    )

    logger.debug("Adding healpix to field coordinates")
    field_coords["field_hpix32"] = lsst.sims.utils.raDec2Hpid(
        32, field_coords.field_ra, field_coords.field_decl
    )

    return field_coords

def _get_sky_mags(sampled_fields_mjd):
    mjd = sampled_fields_mjd.mjd[0]
    assert np.all(mjd == sampled_fields_mjd["mjd"])

    mags = SKY_MODEL.returnMags(mjd, sampled_fields_mjd.field_hpix32, badval=-np.inf)
    sampled_fields_band_mjds = []
    for band in mags:
        sampled_fields_band_mjd = sampled_fields_mjd.copy()
        sampled_fields_band_mjd["band"] = band
        sampled_fields_band_mjd["sky_mag"] = mags[band]
        sampled_fields_band_mjds.append(sampled_fields_band_mjd)
    out_sampled_fields_mjd = pd.concat(sampled_fields_band_mjds)
    return out_sampled_fields_mjd

def _compute_band_fwhms(sampled_fields_band):
    band = sampled_fields_band["band"][0]
    assert np.all(band == sampled_fields_band["band"])

    logger.debug("Calculating the model seeing in %s", band)
    seeing_model = SeeingModel()

    band_idx = ["u", "g", "r", "i", "z", "y"].index(band)
    sampled_fields_band = sampled_fields_band.copy()
    sampled_fields_band["fwhm"] = seeing_model(
        0.7, sampled_fields_band["field_airmass"].values
    )["fwhmEff"][band_idx]
    return sampled_fields_band

def _compute_band_m5(sampled_fields_band):
    band = sampled_fields_band["band"][0]
    assert np.all(band == sampled_fields_band["band"])

    logger.debug("Calculating the 5-sigma magnitude limit in %s", band)
    sampled_fields_band = sampled_fields_band.copy()
    sampled_fields_band["m5"] = lsst.sims.utils.m5_flat_sed(
        band,
        sampled_fields_band["sky_mag"],
        sampled_fields_band["fwhm"],
        30.0,
        sampled_fields_band["field_airmass"],
        nexp=1.0,
    )
    return sampled_fields_band

def read_config(fname):
    """Read m5 configuration file

    Parameters
    ----------
    fname: `str`
        The name of the file to read configuration from.

    Return
    ------
    config: `dict`
        Dictionary of configuration values
    """
    logger.debug("Reading configuration from %s", fname)

    with open(fname, "r") as config_file:
        config = yaml.load(config_file.read(), Loader=yaml.FullLoader)

    # Convert field coordinates into astroy.SkyCoord objects
    for field in config["fields"]:
        config["fields"][field] = SkyCoord(**config["fields"][field])

    return config

def main():
    """Parse command line arguments and config file, and run"""
    parser = ArgumentParser()
    parser.add_argument("config", help="configuration file")
    parser.add_argument("output", help="file in which to write results")

    args = parser.parse_args()
    config_fname = args.config
    output_fname = args.output

    warnings.filterwarnings(
        "ignore",
        message=".*Tried to .* for times after IERS data is valid.*",
        category=astropy.utils.exceptions.AstropyWarning,
    )
    warnings.filterwarnings(
        "ignore",
        message=".*dubious year.*",
        category=astropy.utils.exceptions.AstropyWarning,
    )

    config = read_config(config_fname)

    m5_limits = calc_m5(config)

    (
        m5_limits.set_index(["field_name", "mjd"])
        .drop(columns=["period"])
        .to_hdf(output_fname, "m5")
    )

    return 0

def _init_logger(log_level=logging.DEBUG):
    """Create the ddfpresched logger and set initial configuration"""
    ddfpresched_logger = logging.getLogger("ddfpresched")
    ddfpresched_logger.setLevel(log_level)
    handler = logging.StreamHandler()
    handler.setLevel(log_level)
    formatter = logging.Formatter("%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s")
    handler.setFormatter(formatter)
    ddfpresched_logger.addHandler(handler)
    return ddfpresched_logger

if __name__ == "__main__":
    logger = _init_logger()
    status = main()  # pylint: disable=invalid-name
    sys.exit(status)

import numpy,import astropy,from astropy


from typing import ClassVar, Dict, List, Optional, Sequence, Union, Tuple, Any, Type
from dataclasses import dataclass, field
import numpy as np
from pygmsh.common import bspline
from ..common import Pointer
from pygmsh.geo import Geometry
from math import ceil, floor, sqrt
from geomdl import BSpline
import gmsh  # To do some low level creation

"""
NM: Abbreviation for 'not mine'. Used to denote codes that are took from pyiges
"""

"""
Omitted entity type list
202: Angular Dimension Entity,
210: General Label Entity, 212: General Note Entity, 214: Leader Arrow Entity,
216: Linear Dimension Entity, 218: Ordimate Dimension Entity,
"""

Parameter = Union[str, float]

@dataclass
class Entity:
    type_number: int
    # parameter data pointer
    pd_pointer: int
    structure: int
    line_font_pattern: int
    level: int
    view: int
    transformation_pointer: int
    label_display_associativity: int

    status_number: int
    # instead of using the sequence number, use the index of entity
    # sequence_number: int
    line_weight: int
    color: int
    parameter_line_count: int
    form: int
    entity_label: str
    subscript: int
    # instead of just [], below is required
    # because mutable default is not allowed
    parameters: List[Parameter] = field(default_factory=list)

    _geometry: Optional[Any] = None
    _render: ClassVar[bool] = True

    @staticmethod
    def parse_status_number(x: str) -> int:
        """
        inputs
            x: A candidate status number in string
        output
            status number in integer
        """
        numbers = x.split(" ")
        numbers = [n.zfill(2) for n in numbers]
        return int("".join(numbers))

    def add_parameters(self, *args: List[Parameter]):
        self.parameters += args

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        """Transforms this entity into a vtk object"""
        raise NotImplementedError

    def to_vtk_operator(self) -> Any:
        """
        Transforms this entity into a vtk operator.
        Not an actual object itself, but used to manipulate other vtk objects.
        """
        raise NotImplementedError

@dataclass
class Point(Entity):  # 116

    def add_parameters(self, *args: Tuple[float]):
        self.pos = np.array(args[:3])

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return
        _point = self.pos

        if self.transformation_pointer:
            T = entities[self.transformation_pointer]
            point = T.transform(_point.T)
        else:
            point = _point

        geometry.add_point(point)

@dataclass
class Line(Entity):  # 110

    def add_parameters(self, *args: Tuple[float]):
        self.start = args[:3]
        self.end = args[3:6]

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return

        point1 = np.array(self.start)
        point2 = np.array(self.end)

        if self.transformation_pointer:
            transformation = entities[self.transformation_pointer]
            point1 = transformation.transform(point1)
            point2 = transformation.transform(point2)

        point1 = geometry.add_point(point1, lcar)
        point2 = geometry.add_point(point2, lcar)

        self._geometry = geometry.add_line(point1, point2)

@dataclass
class CopiousData(Entity):  # 106
    """
    Definition of set of points. Type specifies 1: couples, 2: triples, 3: sextuplets

    if type is 1, then 2nd parameter is the common z value for all the other
    xy pairs
    """

    TYPE2PARSE_LENGTH = {1: 2, 2: 3, 3: 6}

    def add_parameters(self, *args: List[Parameter]):
        self.tuple_type: int = int(args[0])

        if self.tuple_type == 1:  # couples
            self.common_z = args[1]
            coordinates = args[2:]
        else:
            self.common_z = None
            coordinates = args[1:]

        length = self.TYPE2PARSE_LENGTH[self.tuple_type]
        self.point_list: List[List] = self.parse_tuple(coordinates, length)

    def parse_tuple(self, args: List[Parameter], length: int) -> List[Tuple]:
        tuple_list: List[List] = []

        # /2 because 2 floats in a couple
        if self.common_z:
            for tuple_idx in range(int(len(args)/length)):
                start = tuple_idx*length
                end = (tuple_idx+1)*length
                tuple_list.append(
                    [*args[start: end], self.common_z])

        else:
            for tuple_idx in range(int(len(args)/length)):
                start = tuple_idx*length
                end = (tuple_idx+1)*length
                tuple_list.append(args[start: end])

        return tuple_list

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return
        """
        Not a mesh, so this can be discarded
        """
        for tuple_point in self.point_list:
            geometry.add_point(tuple_point)

@dataclass
class CircularArc(Entity):  # 100
    """
    Simple circular arc of constant radius. Usually defined with a Transformation Matrix Entity (Type 124).
    """

    def add_parameters(self, *args: List[Parameter]):
        self.z, self.x, self.y, self.x1, self.y1, self.x2, self.y2 = args[:7]

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return

        start = np.array([self.x1, self.y1, self.z])
        end = np.array([self.x2, self.y2, self.z])
        center = np.array([self.x, self.y, self.z])

        if self.transformation_pointer:
            transfomation = entities[self.transformation_pointer]

            start = transfomation.transform(start)[0]
            end = transfomation.transform(end)[0]
            center = transfomation.transform(center)[0]

        start = geometry.add_point(start, lcar)
        end = geometry.add_point(end, lcar)
        center = geometry.add_point(center, lcar)

        geometry.add_circle_arc(start, center, end)

@dataclass
class ConicArc(Entity):  # 124
    """
    Arc defined by the equation: Axt^2 + Bxtyt + Cyt^2 + Dxt + Eyt + F = 0, with a Transformation Matrix (Entity 124). Can define an ellipse, parabola, or hyperbola.

    The definitions of the terms ellipse, parabola, and hyperbola are given in terms of the quantities Q1,Q2, andQ3. These quantities are:

            |  A   B/2  D/2 |        |  A   B/2 |
        Q1= | B/2   C   E/2 |   Q2 = | B/2   C  |   Q3 = A + C
            | D/2  E/2   F  |
    A parent conic curve is:

    An ellipse if Q2 > 0 and Q1Q3 < 0.
    A hyperbola if Q2 < 0 and Q1 != 0.
    A parabola if Q2 = 0 and Q1 != 0.

    +@ from https://en.wikipedia.org/wiki/Matrix_representation_of_conic_sections
    center:
    |xc| = | (BE - 2CD)/(4AC - B^2) |
    |yc| = | (DB - 2AE)/(4AC - B^2) |
    """

    def add_parameters(self, *args: List[Parameter]):
        self.a, self.b, self.c, self.d, self.e, self.f, \
            self.x1, self.y1, self.z1, self.x2, self.y2, self.z2 = args[:12]

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return

        start = np.array([self.x1, self.y1, self.z1])
        end = np.array([self.x2, self.y2, self.z2])
        center = np.array([0, 0, 0])

        a = self.get_major_raidus(
            self.a, self.b, self.c, self.d, self.e, self.f)
        point_on_major = np.array([a, 0, 0])

        if self.transformation_pointer:
            t = entities[self.transformation_pointer]

            start = t.transform(start)[0]
            end = t.transform(end)[0]
            center = t.transform(center)[0]
            point_on_major = t.transform(point_on_major)[0]

        start = geometry.add_point(start, lcar)
        end = geometry.add_point(end, lcar)
        center = geometry.add_point(center, lcar)
        point_on_major = geometry.add_point(point_on_major, lcar)

        geometry.add_ellipse_arc(start, center, point_on_major, end)

    @staticmethod
    def get_x_range(start_pos: np.ndarray, end_pos: np.ndarray, step: float) -> np.ndarray:

        pos = np.arange(start=start_pos[0], stop=end_pos[0]+step, step=step)
        neg = np.array([x * -1 for x in np.flip(pos)])
        return np.hstack((neg, pos))

    @staticmethod
    def get_major_raidus(a, b, c, d, e, f) -> np.float:
        """
        For an ellipse or hyperbola, you can find the canonical form
        from the general form.
        We only need `a`, the major radius so `b` is not calculated
        Reference: https://en.wikipedia.org/wiki/Conic_section#Matrix_notation
        """
        sol = np.array([[a, b/2], [b/2, c]])
        conic_matrix = np.array([[a, b/2, d/2], [b/2, c, e/2], [d/2, e/2, f]])

        eig = np.linalg.eigvals(sol)
        l1, l2 = eig  # lambda 1 and lambda 2
        S = np.linalg.det(conic_matrix)

        canonical_a = -S/(l1**2 * l2)
        return canonical_a

    @staticmethod
    def get_y_positive(a, b, c, d, e, f, x) -> Tuple[float, float]:
        """
        Solve general form for `x`. Given x calculate y the positive version.
        Equations found using `sympy solve`.
        """
        y_positive = (-b*x - d + sqrt(-4*a*c*x**2 - 4*a*e*x -
                                      4*a*f + b**2*x**2 + 2*b*d*x + d**2))/(2*a)

        return y_positive

    @staticmethod
    def get_y_negative(a, b, c, d, e, f, x) -> Tuple[float, float]:
        y_negative = -(b*x + d + sqrt(-4*a*c*x**2 - 4*a*e*x -
                                      4*a*f + b**2*x**2 + 2*b*d*x + d**2))/(2*a)

        return y_negative

@dataclass
class Transformation(Entity):  # 124
    """
    new E = RE + T (E: Entity coordinate)

        | r11 r12 r13 |       | T1 |
    R = | r21 r22 r23 |   T = | T2 |
        | r31 r32 r33 |       | T3 |

    """

    def add_parameters(self, *args: List[float]):
        r11, r12, r13, t1, \
            r21, r22, r23, t2, \
            r31, r32, r33, t3 = (args[:12])

        self.r = np.array([[r11, r12, r13],
                           [r21, r22, r23],
                           [r31, r32, r33]])
        self.t = np.array([[t1, t2, t3]])

    def transform(self, coordinate: np.array) -> np.array:

        return self.r.dot(coordinate) + self.t

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return
        # does nothing by itself
        pass

@dataclass
class SubfigureDefinition(Entity):  # 308

    def add_parameters(self, *args: List[Parameter]):
        self.depth = int(args[0])
        self.name = args[1]
        self.length = int(args[2])
        self.figures = [int(pointer) for pointer in args[3:]]

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return

        if self.transformation_pointer:
            raise NotImplementedError
        pass

@dataclass
class SingularSubfigureInstance(Entity):  # 408

    def add_parameters(self, *args: List[Parameter]):
        self.pointer = int(args[0])
        self.translation = np.array(args[1:4])
        self.scale_factor = args[4]

@dataclass
class RationalBSplineCurve(Entity):
    """
    Composes analytic curves.
    Form: 0=Determined by data 1=Line 2=Circular arc 3=Eliptic arc 4=Parabolic arc 5=Hyperbolic arc
    """

    def add_parameters(self, *args: List[Parameter]):
        self.K, self.degree, self.flag1, self.flag2, self.flag3, self.flag4 \
            = map(lambda x: int(x), args[:6])

        weight_start_idx = 8 + self.K + self.degree
        control_start_idx = weight_start_idx + self.K + 1
        etc_start_idx = control_start_idx + 3 * self.K + 3
        knot_range = range(6, weight_start_idx)
        weight_range = range(weight_start_idx, control_start_idx)

        self.v0, self.v1, self.xn, self.yn, self.zn = args[etc_start_idx:]

        self.knots = [args[i] for i in knot_range]
        self.weights = [args[i] for i in weight_range]
        temp_control_points = args[control_start_idx:etc_start_idx]
        self.control_points = [temp_control_points[i:i + 3] for i in range(0, len(temp_control_points), 3)]

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return

        curve = BSpline.Curve()
        curve.degree = self.degree
        curve.ctrlpts = self.control_points
        curve.knotvector = self.knots
        curve.delta = lcar

        _points = curve.evalpts

        points = []
        for point in _points:

            if self.transformation_pointer:
                transfomation = entities[self.transformation_pointer]

                point = transfomation.transform(point)[0]

            point = geometry.add_point(point, lcar)
            points.append(point)

        spline = geometry.add_bspline(points)

@dataclass
class RationalBSplineSurface(Entity):
    """
    This is a surface entity defined by multiple surfaces. The form number describes the general type: 0=determined from data, 1=Plane, 2=Right circular cylinder, 3=Cone, 4=Sphere, 5=Torus, 6=Surface of revolution, 7=Tabulated cylinder, 8=Ruled surface, 9=General quadratic surface.

    Takes up most of the time saving.
    """

    def add_parameters(self, *args: List[Parameter]):
        start_idx = 9
        self.k1, self.k2, self.m1, self.m2, \
            self.flag1, self.flag2, self.flag3, self.flag4, self.flag5 \
            = map(lambda x: int(x), args[:start_idx])

        knot2_start = start_idx + self.k2 + self.m1 + 2
        weight_start = knot2_start + self.k1 + self.m2 + 2
        self.knot1 = args[start_idx: knot2_start]
        self.knot2 = args[knot2_start: weight_start]

        weight_count = (self.k2 + 1) * (self.k1 + 1)
        control_start = weight_start + weight_count
        self.weights = args[weight_start: control_start]

        etc_start = control_start + 3 * weight_count

        self.control_points: List[List] = []

        i = control_start
        while i < etc_start:
            self.control_points.append(np.array(args[i: i+3]))
            i += 3

        self.control_points = np.array(
            self.control_points).reshape(self.k2+1, self.k1 + 1, 3)

        self.U0, self.U1, self.V0, self.V1 = args[etc_start:]

        assert len(self.knot1) == (2 + self.k2 + self.m1)
        assert len(self.knot2) == (2 + self.k1 + self.m2)
        assert len(self.weights) == weight_count
        assert self.control_points.size == (weight_count * 3)

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return

        rows, cols, z = self.control_points.shape
        # The id have to match, not just coordinates
        top = []
        for i in range(0, cols):
            point = self.control_points[0, i]
            top.append(geometry.add_point(point, lcar))
        top_spline = geometry.add_bspline(top)

        right = [top[-1]]
        for i in range(1, rows):
            point = self.control_points[i, cols-1]
            right.append(geometry.add_point(point, lcar))
        right_spline = geometry.add_bspline(right)

        bottom = [right[-1]]
        for i in range(cols-2, -1, -1):
            point = self.control_points[rows-1, i]
            bottom.append(geometry.add_point(point, lcar))
        bottom_spline = geometry.add_bspline(bottom)

        left = [bottom[-1]]
        for i in range(rows-2, 0, -1):
            point = self.control_points[i, 0]
            left.append(geometry.add_point(point, lcar))
        left.append(top[0])
        left_spline = geometry.add_bspline(left)

        loop = geometry.add_curve_loop(
            [top_spline, right_spline, bottom_spline, left_spline])
        surface = geometry.add_surface(loop)

@dataclass
class Face(Entity):
    """
    Defines a bound portion of three dimensional space (R^3) which has a finite area. Used to construct B-Rep Geometries
    """

    def add_parameters(self, *args: List[Parameter]):
        self.pointer = int(args[0])
        self.loop_count = int(args[1])
        self.loop_flag = int(args[2])

        self.loops = [args[i] for i in range(3, 2 + self.loop_count)]

@dataclass
class LoopEdge:
    """Represents edge in a loop. NOT AN ENTITY"""

    type: int
    pointer: int
    index: int
    flag: bool
    curves: List[Tuple[bool, int]]

@dataclass
class Loop(Entity):
    """Defines a loop, specifying a bounded face, for B-Rep Geometries."""

    def add_parameters(self, *args: List[Parameter]):
        self.edge_count = int(args[0])
        self.edges: List[LoopEdge] = []

        edge_idx = 0
        arg_idx = 1
        while edge_idx < self.edge_count:
            edge_type = int(args[arg_idx])
            edge_list_pointer = int(args[arg_idx + 1])
            edge_index = int(args[arg_idx + 2])
            flag = bool(args[arg_idx + 3])
            curve_count = int(args[arg_idx + 4])
            arg_idx += 5

            curves = []
            for i in range(curve_count):
                iso = bool(args[arg_idx])
                psc = int(args[arg_idx + 1])
                curves.append((iso, psc))
                arg_idx += 2

            edge = LoopEdge(edge_type, edge_list_pointer,
                            edge_index, flag, curves)
            self.edges.append(edge)

            edge_idx += 1

@dataclass
class Edge:
    curve: int
    start_vertex_list: int
    start_vertex_index: int
    end_vertex_list: int
    end_vertex_index: int

@dataclass
class EdgeList(Entity):

    def add_parameters(self, *args: List[Parameter]):
        edge_count = int(args[0])

        edge_index = 0
        arg_index = 1
        self.edges: List[Edge] = []

        while edge_index < edge_count:
            e = Edge(*map(lambda x: int(x), args[arg_index:arg_index+5]))
            self.edges.append(e)

            arg_index += 5
            edge_index += 1

@dataclass
class Vertex:
    x: float
    y: float
    z: float

@dataclass
class VertexList(Entity):

    def add_parameters(self, *args: List[Parameter]):
        vertex_count = int(args[0])

        vertex_index = 0
        arg_index = 1
        self.vertices: List[Vertex] = []

        while vertex_index < vertex_count:
            v = Vertex(*args[arg_index: arg_index + 3])
            self.vertices.append(v)
            arg_index += 3
            vertex_index += 1

@dataclass
class SurfaceOfRevolution(Entity):  # 120

    def add_parameters(self, *args: List[Parameter]):
        self.axis_line = Pointer(args[0])
        self.target = Pointer(args[1])
        self.start_angle, self.end_angle = args[2:4]  # radian

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return
        line = entities[self.axis_line]
        target = entities[self.target]
        if not target._geometry:
            # this entity is dependent on another entity,
            # if the other entity is not rendered, this will not work
            # so enable it for this
            old_render = target._render
            target._render = True
            target.to_vtk(entities, geometry, lcar)
            target._render = old_render
        target_geometry = target._geometry

        geometry.revolve(target_geometry, line.start, line.end, 3.14)

@dataclass
class CompositeCurve(Entity):  # 102

    def add_parameters(self, *args: List[Parameter]):
        curve_count = int(args[0])
        self.curves: List[Pointer] = list(map(Pointer, args[1:]))

    def to_vtk(self, entities: Dict[Pointer, "Entity"], geometry: Geometry, lcar: float = 0.1):
        if not self._render:
            return

# todo
# Surface Revolution Type 120
# Composite Curve 102
# Trimmed Surface 144
# Curve on Parametric Surface 142