# param_sweep.py
#
# This file is part of scqubits: a Python package for superconducting qubits,
# arXiv:2107.08552 (2021). https://arxiv.org/abs/2107.08552
#
#    Copyright (c) 2019 and later, Jens Koch and Peter Groszkowski
#    All rights reserved.
#
#    This source code is licensed under the BSD-style license found in the
#    LICENSE file in the root directory of this source tree.
############################################################################

import functools
import weakref

from abc import ABC
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np

from numpy import ndarray
from qutip.qobj import Qobj

import scqubits.core.central_dispatch as dispatch
import scqubits.core.descriptors as descriptors
import scqubits.core.hilbert_space as hspace
import scqubits.core.spec_lookup as spec_lookup
import scqubits.core.storage as storage
import scqubits.io_utils.fileio_qutip as qutip_serializer
import scqubits.io_utils.fileio_serializers as serializers
import scqubits.settings as settings
import scqubits.utils.cpu_switch as cpu_switch
import scqubits.utils.misc as utils

from scqubits.core.hilbert_space import HilbertSpace
from scqubits.core.oscillator import Oscillator
from scqubits.core.qubit_base import QubitBaseClass
from scqubits.core.spec_lookup import SpectrumLookup
from scqubits.core.storage import DataStore, SpectrumData
from scqubits.io_utils.fileio_qutip import QutipEigenstates

if TYPE_CHECKING:
    from scqubits.io_utils.fileio import IOData

if settings.IN_IPYTHON:
    from tqdm.notebook import tqdm
else:
    from tqdm import tqdm

from scqubits.utils.typedefs import QuantumSys

class _ParameterSweepBase(ABC):
    """
    The _ParameterSweepBase class is an abstract base class for ParameterSweep and
    StoredSweep
    """

    param_name = descriptors.WatchedProperty(str, "PARAMETERSWEEP_UPDATE")
    param_vals = descriptors.WatchedProperty(ndarray, "PARAMETERSWEEP_UPDATE")
    param_count = descriptors.WatchedProperty(int, "PARAMETERSWEEP_UPDATE")
    evals_count = descriptors.WatchedProperty(int, "PARAMETERSWEEP_UPDATE")
    lookup = descriptors.ReadOnlyProperty(SpectrumLookup)
    _hilbertspace: hspace.HilbertSpace

    def get_subsys(self, index: int) -> QuantumSys:
        return self._hilbertspace[index]

    def get_subsys_index(self, subsys: QuantumSys) -> int:
        return self._hilbertspace.get_subsys_index(subsys)

    @property
    def osc_subsys_list(self) -> List[Oscillator]:
        return self._hilbertspace.osc_subsys_list

    @property
    def qbt_subsys_list(self) -> List[QubitBaseClass]:
        return self._hilbertspace.qbt_subsys_list

    @property
    def subsystem_count(self) -> int:
        return self._hilbertspace.subsystem_count

    @property
    def bare_specdata_list(self) -> List[SpectrumData]:
        return self.lookup._bare_specdata_list

    @property
    def dressed_specdata(self) -> SpectrumData:
        return self.lookup._dressed_specdata

    def _lookup_bare_eigenstates(
        self,
        param_index: int,
        subsys: QuantumSys,
        bare_specdata_list: List[SpectrumData],
    ) -> Union[ndarray, List[QutipEigenstates]]:
        """
        Parameters
        ----------
        param_index:
            position index of parameter value in question
        subsys:
            Hilbert space subsystem for which bare eigendata is to be looked up
        bare_specdata_list:
            may be provided during partial generation of the lookup

        Returns
        -------
            bare eigenvectors for the specified subsystem and the external parameter
            fixed to the value indicated by its index
        """
        subsys_index = self.get_subsys_index(subsys)
        return bare_specdata_list[subsys_index].state_table[param_index]  # type: ignore

    @property
    def system_params(self) -> Dict[str, Any]:
        return self._hilbertspace.get_initdata()

    def new_datastore(self, **kwargs) -> DataStore:
        """Return DataStore object with system/sweep information obtained from self."""
        return storage.DataStore(
            self.system_params, self.param_name, self.param_vals, **kwargs
        )

class _ParameterSweep(
    _ParameterSweepBase, dispatch.DispatchClient, serializers.Serializable
):
    """
    The ParameterSweep class helps generate spectral and associated data for a
    composite quantum system, as an externa, parameter, such as flux, is swept over
    some given interval of values. Upon initialization, these data are calculated and
    stored internally, so that plots can be generated efficiently. This is of
    particular use for interactive displays used in the Explorer_ class.

    Parameters
    ----------
    param_name:
        name of external parameter to be varied
    param_vals:
        array of parameter values
    evals_count:
        number of eigenvalues and eigenstates to be calculated for the composite
        Hilbert space
    hilbertspace:
        collects all data specifying the Hilbert space of interest
    subsys_update_list:
        list of subsys_list in the Hilbert space which get modified when the external
        parameter changes
    update_hilbertspace:
        update_hilbertspace(param_val) specifies how a change in the external
        parameter affects the Hilbert space components
    num_cpus:
        number of CPUS requested for computing the sweep (default: settings.NUM_CPUS)
    """

    param_name = descriptors.WatchedProperty(str, "PARAMETERSWEEP_UPDATE")
    param_vals = descriptors.WatchedProperty(ndarray, "PARAMETERSWEEP_UPDATE")
    param_count = descriptors.WatchedProperty(int, "PARAMETERSWEEP_UPDATE")
    evals_count = descriptors.WatchedProperty(int, "PARAMETERSWEEP_UPDATE")
    subsys_update_list = descriptors.WatchedProperty(
        List[QuantumSys], "PARAMETERSWEEP_UPDATE"
    )
    update_hilbertspace = descriptors.WatchedProperty(Callable, "PARAMETERSWEEP_UPDATE")
    lookup = descriptors.ReadOnlyProperty(SpectrumLookup)

    def __init__(
        self,
        param_name: str,
        param_vals: ndarray,
        evals_count: int,
        hilbertspace: HilbertSpace,
        subsys_update_list: List[QuantumSys],
        update_hilbertspace: Callable,
        num_cpus: Optional[int] = None,
    ) -> None:
        num_cpus = num_cpus or settings.NUM_CPUS
        self.param_name = param_name
        self.param_vals = param_vals
        self.param_count = len(param_vals)
        self.evals_count = evals_count
        self._hilbertspace = hilbertspace
        self.subsys_update_list = tuple(subsys_update_list)
        self.update_hilbertspace = update_hilbertspace
        self.num_cpus = num_cpus
        self._lookup: Union[SpectrumLookup, None] = None
        self._bare_hamiltonian_constant: Qobj

        self.tqdm_disabled = settings.PROGRESSBAR_DISABLED or (num_cpus > 1)

        dispatch.CENTRAL_DISPATCH.register("PARAMETERSWEEP_UPDATE", self)
        dispatch.CENTRAL_DISPATCH.register("HILBERTSPACE_UPDATE", self)

        # generate the spectral data sweep
        if settings.AUTORUN_SWEEP:
            self.run()

    def run(self) -> None:
        """Top-level method for generating all parameter sweep data"""
        self.cause_dispatch()  # one dispatch before temp. disabling CENTRAL_DISPATCH
        settings.DISPATCH_ENABLED = False
        bare_specdata_list = self._compute_bare_specdata_sweep()
        dressed_specdata = self._compute_dressed_specdata_sweep(bare_specdata_list)
        self._lookup = spec_lookup.SpectrumLookup(
            self, dressed_specdata, bare_specdata_list
        )
        settings.DISPATCH_ENABLED = True

    # HilbertSpace: methods for CentralDispatch ---------------------------------------
    def cause_dispatch(self) -> None:
        self.update_hilbertspace(self.param_vals[0])

    def receive(self, event: str, sender: object, **kwargs) -> None:
        """Hook to CENTRAL_DISPATCH. This method is accessed by the global
        CentralDispatch instance whenever an event occurs that ParameterSweep is
        registered for. In reaction to update events, the lookup table is marked as
        out of sync.

        Parameters
        ----------
        event:
            type of event being received
        sender:
            identity of sender announcing the event
        **kwargs
        """
        if self._lookup is not None:
            if event == "HILBERTSPACE_UPDATE" and sender is self._hilbertspace:
                self._lookup._out_of_sync = True
                # print('Lookup table now out of sync')
            elif event == "PARAMETERSWEEP_UPDATE" and sender is self:
                self._lookup._out_of_sync = True
                # print('Lookup table now out of sync')

    # ParameterSweep: file IO methods -------------------------------------------------
    @classmethod
    def deserialize(cls, iodata: "IOData") -> "_StoredSweep":
        """
        Take the given IOData and return an instance of the described class,
        initialized with the data stored in io_data.

        Parameters
        ----------
        iodata: IOData

        Returns
        -------
        _StoredSweep
        """
        data_dict = iodata.as_kwargs()
        lookup = data_dict.pop("_lookup")
        data_dict["dressed_specdata"] = lookup._dressed_specdata
        data_dict["bare_specdata_list"] = lookup._bare_specdata_list
        new_storedsweep = _StoredSweep(**data_dict)
        new_storedsweep._lookup = lookup
        return new_storedsweep

    def serialize(self) -> "IOData":
        """
        Convert the content of the current class instance into IOData format.

        Returns
        -------
        IOData
        """
        if self._lookup is None:
            raise ValueError("Nothing to save - no lookup data has been generated yet.")

        initdata = {
            "param_name": self.param_name,
            "param_vals": self.param_vals,
            "evals_count": self.evals_count,
            "hilbertspace": self._hilbertspace,
            "_lookup": self._lookup,
        }
        iodata = serializers.dict_serialize(initdata)
        iodata.typename = "_StoredSweep"
        return iodata

    # ParameterSweep: private methods for generating the sweep ------------------------
    def _compute_bare_specdata_sweep(self) -> List[SpectrumData]:
        """
        Pre-calculates all bare spectral data needed for the interactive explorer
        display.
        """
        bare_eigendata_constant = [
            self._compute_bare_spectrum_constant()
        ] * self.param_count
        target_map = cpu_switch.get_map_method(self.num_cpus)
        with utils.InfoBar(
            "Parallel compute bare eigensys [num_cpus={}]".format(self.num_cpus),
            self.num_cpus,
        ):
            bare_eigendata_varying = list(
                target_map(
                    self._compute_bare_spectrum_varying,
                    tqdm(
                        self.param_vals,
                        desc="Bare spectra",
                        leave=False,
                        disable=self.tqdm_disabled,
                    ),
                )
            )
        bare_specdata_list = self._recast_bare_eigendata(
            bare_eigendata_constant, bare_eigendata_varying
        )
        del bare_eigendata_constant
        del bare_eigendata_varying
        return bare_specdata_list

    def _compute_dressed_specdata_sweep(
        self, bare_specdata_list: List[SpectrumData]
    ) -> SpectrumData:
        """
        Calculates and returns all dressed spectral data.
        """
        self._bare_hamiltonian_constant = self._compute_bare_hamiltonian_constant(
            bare_specdata_list
        )
        param_indices = range(self.param_count)
        func = functools.partial(
            self._compute_dressed_eigensystem, bare_specdata_list=bare_specdata_list
        )
        target_map = cpu_switch.get_map_method(self.num_cpus)

        with utils.InfoBar(
            "Parallel compute dressed eigensys [num_cpus={}]".format(self.num_cpus),
            self.num_cpus,
        ):
            dressed_eigendata = list(
                target_map(
                    func,
                    tqdm(
                        param_indices,
                        desc="Dressed spectrum",
                        leave=False,
                        disable=self.tqdm_disabled,
                    ),
                )
            )
        dressed_specdata = self._recast_dressed_eigendata(dressed_eigendata)
        del dressed_eigendata
        return dressed_specdata

    def _recast_bare_eigendata(
        self,
        static_eigendata: List[List[Tuple[ndarray, ndarray]]],
        bare_eigendata: List[List[Tuple[ndarray, ndarray]]],
    ) -> List[SpectrumData]:
        specdata_list = []
        for index, subsys in enumerate(self._hilbertspace):
            if subsys in self.subsys_update_list:
                eigendata = bare_eigendata
            else:
                eigendata = static_eigendata
            evals_count = subsys.truncated_dim
            dim = subsys.hilbertdim()
            esys_dtype = subsys._evec_dtype

            energy_table = np.empty(
                shape=(self.param_count, evals_count), dtype=np.float_
            )
            state_table = np.empty(
                shape=(self.param_count, dim, evals_count), dtype=esys_dtype
            )
            for j in range(self.param_count):
                energy_table[j] = eigendata[j][index][0]
                state_table[j] = eigendata[j][index][1]
            specdata_list.append(
                storage.SpectrumData(
                    energy_table,
                    system_params={},
                    param_name=self.param_name,
                    param_vals=self.param_vals,
                    state_table=state_table,
                )
            )
        return specdata_list

    def _recast_dressed_eigendata(
        self, dressed_eigendata: List[Tuple[ndarray, QutipEigenstates]]
    ) -> SpectrumData:
        evals_count = self.evals_count
        energy_table = np.empty(shape=(self.param_count, evals_count), dtype=np.float_)
        state_table = []  # for dressed states, entries are Qobj
        for j in range(self.param_count):
            energy_table[j] = np.real_if_close(dressed_eigendata[j][0])
            state_table.append(dressed_eigendata[j][1])
        specdata = storage.SpectrumData(
            energy_table,
            system_params={},
            param_name=self.param_name,
            param_vals=self.param_vals,
            state_table=state_table,
        )
        return specdata

    def _compute_bare_hamiltonian_constant(
        self, bare_specdata_list: List[SpectrumData]
    ) -> Qobj:
        """
        Returns
        -------
            composite Hamiltonian composed of bare Hamiltonians of subsys_list
            independent of the external parameter
        """
        static_hamiltonian = 0
        for index, subsys in enumerate(self._hilbertspace):
            if subsys not in self.subsys_update_list:
                evals = bare_specdata_list[index].energy_table[0]
                static_hamiltonian += self._hilbertspace.diag_hamiltonian(subsys, evals)
        return static_hamiltonian

    def _compute_bare_hamiltonian_varying(
        self, bare_specdata_list: List[SpectrumData], param_index: int
    ) -> Qobj:
        """
        Parameters
        ----------
        param_index:
            position index of current value of the external parameter

        Returns
        -------
            composite Hamiltonian consisting of all bare Hamiltonians which depend on
            the external parameter
        """
        hamiltonian = 0
        for index, subsys in enumerate(self._hilbertspace):
            if subsys in self.subsys_update_list:
                evals = bare_specdata_list[index].energy_table[param_index]
                hamiltonian += self._hilbertspace.diag_hamiltonian(subsys, evals)
        return hamiltonian

    def _compute_bare_spectrum_constant(self) -> List[Tuple[ndarray, ndarray]]:
        """
        Returns
        -------
            eigensystem data for each subsystem that is not affected by a change of the
            external parameter
        """
        eigendata = []
        for subsys in self._hilbertspace:
            if subsys not in self.subsys_update_list:
                evals_count = subsys.truncated_dim
                eigendata.append(subsys.eigensys(evals_count=evals_count))
            else:
                eigendata.append(None)  # type: ignore
        return eigendata

    def _compute_bare_spectrum_varying(
        self, param_val: float
    ) -> List[Tuple[ndarray, ndarray]]:
        """
        For given external parameter value obtain the bare eigenspectra of each bare
        subsystem that is affected by changes in the external parameter. Formulated
        to be used with Pool.map()

        Returns
        -------
            (evals, evecs) bare eigendata for each subsystem that is parameter-dependent
        """
        eigendata = []
        self.update_hilbertspace(param_val)
        for subsys in self._hilbertspace:
            if subsys in self.subsys_update_list:
                evals_count = subsys.truncated_dim
                subsys_index = self._hilbertspace.get_subsys_index(subsys)
                eigendata.append(
                    self._hilbertspace[subsys_index].eigensys(evals_count=evals_count)
                )
            else:
                eigendata.append(None)  # type: ignore
        return eigendata

    def _compute_dressed_eigensystem(
        self, param_index: int, bare_specdata_list: List[SpectrumData]
    ) -> Tuple[ndarray, QutipEigenstates]:
        hamiltonian = (
            self._bare_hamiltonian_constant
            + self._compute_bare_hamiltonian_varying(bare_specdata_list, param_index)
        )

        for interaction_term in self._hilbertspace.interaction_list:
            evecs1 = self._lookup_bare_eigenstates(
                param_index, interaction_term.subsys1, bare_specdata_list
            )
            evecs2 = self._lookup_bare_eigenstates(
                param_index, interaction_term.subsys2, bare_specdata_list
            )
            hamiltonian += self._hilbertspace.interactionterm_hamiltonian(
                interaction_term, evecs1=evecs1, evecs2=evecs2
            )
        evals, evecs = hamiltonian.eigenstates(eigvals=self.evals_count)
        evecs = evecs.view(qutip_serializer.QutipEigenstates)
        return evals, evecs

    def _lookup_bare_eigenstates(
        self,
        param_index: int,
        subsys: QuantumSys,
        bare_specdata_list: List[SpectrumData],
    ) -> ndarray:
        """
        Parameters
        ----------
        param_index:
            position index of parameter value in question
        subsys:
            Hilbert space subsystem for which bare eigendata is to be looked up
        bare_specdata_list:
            may be provided during partial generation of the lookup

        Returns
        -------
            bare eigenvectors for the specified subsystem and the external parameter
            fixed to the value indicated by its index
        """
        subsys_index = self.get_subsys_index(subsys)
        return bare_specdata_list[subsys_index].state_table[param_index]  # type: ignore

class _StoredSweep(
    _ParameterSweepBase, dispatch.DispatchClient, serializers.Serializable
):
    param_name = descriptors.WatchedProperty(str, "PARAMETERSWEEP_UPDATE")
    param_vals = descriptors.WatchedProperty(ndarray, "PARAMETERSWEEP_UPDATE")
    param_count = descriptors.WatchedProperty(int, "PARAMETERSWEEP_UPDATE")
    evals_count = descriptors.WatchedProperty(int, "PARAMETERSWEEP_UPDATE")
    lookup = descriptors.ReadOnlyProperty(SpectrumLookup)

    def __init__(
        self,
        param_name: str,
        param_vals: ndarray,
        evals_count: int,
        hilbertspace: HilbertSpace,
        dressed_specdata: SpectrumData,
        bare_specdata_list: List[SpectrumData],
    ) -> None:
        self.param_name = param_name
        self.param_vals = param_vals
        self.param_count = len(param_vals)
        self.evals_count = evals_count
        self._hilbertspace = hilbertspace
        self._lookup = spec_lookup.SpectrumLookup(
            hilbertspace, dressed_specdata, bare_specdata_list, auto_run=False
        )

    # StoredSweep: file IO methods -----------------------------------------------------
    @classmethod
    def deserialize(cls, iodata: "IOData") -> "_StoredSweep":
        """
        Take the given IOData and return an instance of the described class,
        initialized with the data stored in io_data.

        Parameters
        ----------
        iodata: IOData

        Returns
        -------
        _StoredSweep
        """
        data_dict = iodata.as_kwargs()
        lookup = data_dict.pop("_lookup")
        data_dict["dressed_specdata"] = lookup._dressed_specdata
        data_dict["bare_specdata_list"] = lookup._bare_specdata_list
        new_storedsweep = _StoredSweep(**data_dict)
        new_storedsweep._lookup = lookup
        new_storedsweep._lookup._hilbertspace = weakref.proxy(
            new_storedsweep._hilbertspace
        )
        return new_storedsweep

    # _StoredSweep: other methods
    def get_hilbertspace(self) -> HilbertSpace:
        return self._hilbertspace

    def new_sweep(
        self,
        subsys_update_list: List[QuantumSys],
        update_hilbertspace: Callable,
        num_cpus: Optional[int] = None,
    ) -> _ParameterSweep:
        num_cpus = num_cpus or settings.NUM_CPUS
        return _ParameterSweep(
            self.param_name,
            self.param_vals,
            self.evals_count,
            hilbertspace=self._hilbertspace,
            subsys_update_list=subsys_update_list,
            update_hilbertspace=update_hilbertspace,
            num_cpus=num_cpus,
        )


#coding:utf-8
import numpy as np
import tensorflow as tf
from .Model import Model

if tf.__version__ > '0.12.1':
	matmul_func = tf.matmul
else:
	matmul_func = tf.batch_matmul

class TransNB(Model):
	r'''
	TransE is the first model to introduce translation-based embedding,
	which interprets relations as the translations operating on entities.
	'''
	def _transfer(self, transfer_matrix, embeddings):
		return matmul_func(embeddings, transfer_matrix)

	def _calc(self, h, t, r):
		h = tf.nn.l2_normalize(h, -1)
		t = tf.nn.l2_normalize(t, -1)
		r = tf.nn.l2_normalize(r, -1)
		return abs(h + r - t)

	def embedding_def(self):
		#Obtaining the initial configuration of the model
		config = self.get_config()
		#Defining required parameters of the model, including embeddings of entities and relations
		self.ent_embeddings = tf.get_variable(name = "ent_embeddings", shape = [config.entTotal, config.hidden_size], initializer = tf.contrib.layers.xavier_initializer(uniform = False))
		self.rel_embeddings = tf.get_variable(name = "rel_embeddings", shape = [config.relTotal, config.hidden_size], initializer = tf.contrib.layers.xavier_initializer(uniform = False))
        self.type_transfer_matrix = tf.get_variable(name = "type_transfer_matrix", shape = [config.rel_type, config.ent_size * config.rel_size], initializer = tf.contrib.layers.xavier_initializer(uniform = False))
        self.parameter_lists = {"ent_embeddings":self.ent_embeddings, \
								"rel_embeddings":self.rel_embeddings,
                                "type_transfer_matrix":self.type_transfer_matrix}

	def loss_def(self):
		#Obtaining the initial configuration of the model
		config = self.get_config()
		#To get positive triples and negative triples for training
		#The shapes of pos_h, pos_t, pos_r are (batch_size, 1)
		#The shapes of neg_h, neg_t, neg_r are (batch_size, negative_ent + negative_rel)
		pos_h, pos_t, pos_r, pos_type_r = self.get_positive_instance(in_batch = True)
		neg_h, neg_t, neg_r, neg_type_r = self.get_negative_instance(in_batch = True)
		#Embedding entities and relations of triples, e.g. p_h, p_t and p_r are embeddings for positive triples
		self.p_h = tf.nn.embedding_lookup(self.ent_embeddings, pos_h)
		self.p_t = tf.nn.embedding_lookup(self.ent_embeddings, pos_t)
		self.p_r = tf.nn.embedding_lookup(self.rel_embeddings, pos_r)
        self.p_type_r = tf.reshape(tf.nn.embedding_lookup(self.type_transfer_matrix, pos_type_r), [-1, config.ent_size, config.rel_size])

		self.n_h = tf.nn.embedding_lookup(self.ent_embeddings, neg_h)
		self.n_t = tf.nn.embedding_lookup(self.ent_embeddings, neg_t)
		self.n_r = tf.nn.embedding_lookup(self.rel_embeddings, neg_r)
        self.n_type_r = tf.reshape(tf.nn.embedding_lookup(self.type_transfer_matrix, neg_type_r), [-1, config.ent_size, config.rel_size])

        self.p_h_ = self._transfer(self.p_type_r, self.p_h)
        self.p_t_ = self._transfer(self.p_type_r, self.p_t)
        self.n_h_ = self._transfer(self.n_type_r, self.n_h)
        self.n_t_ = self._transfer(self.n_type_r, self.n_t)
		#Calculating score functions for all positive triples and negative triples
		#The shape of _p_score is (batch_size, 1, hidden_size)
		#The shape of _n_score is (batch_size, negative_ent + negative_rel, hidden_size)
		self._p_score = self._calc(self.p_h_, self.p_t_, self.p_r)
		self._n_score = self._calc(self.n_h_, self.n_t_, self.n_r)
		#The shape of p_score is (batch_size, 1, 1)
		#The shape of n_score is (batch_size, negative_ent + negative_rel, 1)
		self.p_score =  tf.reduce_sum(self._p_score, -1, keep_dims = True)
		self.n_score =  tf.reduce_sum(self._n_score, -1, keep_dims = True)
		#Calculating loss to get what the framework will optimize
		self.loss = tf.reduce_mean(tf.maximum(self.p_score - self.n_score + config.margin, 0))

	def predict_def(self):
		predict_h, predict_t, predict_r = self.get_predict_instance()
		predict_h_e = tf.nn.embedding_lookup(self.ent_embeddings, predict_h)
		predict_t_e = tf.nn.embedding_lookup(self.ent_embeddings, predict_t)
		predict_r_e = tf.nn.embedding_lookup(self.rel_embeddings, predict_r)
		self.predict = tf.reduce_mean(self._calc(predict_h_e, predict_t_e, predict_r_e), 1, keep_dims = False)


#! /usr/bin/env python
"""
This module takes a table of exposures, finds all exposures tagged as planetary nebula
exposures, and:

- For each grism
    - For each spectral order
        - For each file
            - Determines which emission lines fall in that order of that file
            - Attempts to automatically fit the line location
            - Allows the user to override the result of the automatic fit
            - Records the resulting fit location and status (good/bad/custom/etc.)

The module then saves a table of the results.

In addition, the module can take the result table above and use it to derive an overall
wavelength fit for each grism/order value of the inputs, and produce yet another output
table given the results of that fit.

Authors
-------
- Brian York (all python code)
- Ralph Bohlin (original IDL code)

Use
---
This module can be run from the command line (although one of the `abscal.commands` or
`abscal.idl_commands` scripts would be preferred for that), but is mostly intended to be
imported, either by binary scripts or for use from within python::

    from abscal.wfc3.reduce_grism_wavelength import wlmeas, wlmake

    interim_table = wlmeas(input_table, command_line_arg_namespace, override_dict)
    final_table = wlmake(input_table, interim_table, command_line_arg_namespace, override_dict)

The override dict allows for many of the default input parameters to be overriden (as
defaults -- individual per-exposure overrides defined in the data files will still take
priority). There are currently no default parameters that can be overriden in this module.
"""

# *****TODO*****
#   The following things could potentially be overridable parameters:
#       - wlmeas
#           - search region size around emission line
#           - order wavelength search ranges (wrang)

import datetime
import glob
import json
import os
import yaml

import matplotlib.pyplot as plt
import numpy as np

from astropy import constants as consts
from astropy.io import ascii, fits
from astropy.table import Table, Column, unique
from astropy.time import Time
from copy import deepcopy
from matplotlib.widgets import TextBox
from pathlib import Path
from photutils.detection import DAOStarFinder
from scipy.linalg import lstsq
from scipy.stats import mode

from abscal.common.args import parse
from abscal.common.standard_stars import find_star_by_name
from abscal.common.utils import air2vac, get_data_file, get_defaults, linecen
from abscal.common.utils import smooth_model, tabinv
from abscal.common.exposure_data_table import AbscalDataTable
from abscal.wfc3.reduce_grism_extract import reduce

def wlimaz(root, y_arr, wave_arr, directory, verbose):
    """
    Find a line from the IMA zero-read.

    If the very bright 10380A line falls in the first order, you can end up
    trying to centre a saturated line. In this case, use the _ima.fits file,
    which holds all of the individual reads, and measure the line centre from
    the zero-read ima file in the first order.

    Parameters
    ----------
    root : str
        The file name to be checked
    y_arr : np.ndarray
        The y-values (flux values) from the flt file
    wave_arr : np.ndarray
        The approximate wavelength values from the flt file
    directory : str
        The directory where the flt file is located (and where the ima file
        should be located)

    Returns
    -------
    star_x : float
        The x centre of the line
    star_y : float
        The y centre of the line
    """
    file_name = os.path.join(directory, root+"_ima.fits")
    with fits.open(file_name) as in_file:
        hd = in_file[0].header
        nexten = hd['NEXTEND']
        ima = in_file[nexten-4].data    # zero read
        ima = ima[5:1019,5:1019]        # trim to match flt
        dq = in_file[nexten-2].data     # DQ=8 is unstable in Zread
        dq = dq[5:1019,5:1019]          # trim to match flt

        # Get approximate position from preliminary extraction
        xlin = tabinv(wave_arr, np.array((10830.,)))
        xapprox = np.floor(xlin + .5).astype(np.int32)
        if isinstance(xapprox, np.ndarray):
            xapprox = xapprox[0]
        yapprox = np.floor(y_arr[xapprox] + .5).astype(np.int32)
        if isinstance(yapprox, np.ndarray):
            yapprox = yapprox[0]
        if verbose:
            print(type(xapprox), type(yapprox))
            msg = "WLIMAZ: 10830 line at approx ({},{})"
            print(msg.format(xapprox, yapprox))

        # Fix any DQ=8 pixels
        ns = 11 # for an 11x11 search area
        sbimg = ima[yapprox-ns//2:yapprox+ns//2+1,xapprox-ns//2:xapprox+ns//2+1]
        sbdq = dq[yapprox-ns//2:yapprox+ns//2+1,xapprox-ns//2:xapprox+ns//2+1]
        bad = np.where((sbdq & 8) != 0)
        if len(bad) > 0:
            nbad = len(bad[0])
        # from wlimaz.pro comment:
        # ; I see up to nbad=5 in later data, eg ic6906bzq, but ima NOT zeroed & looks OK
        # ; no response from SED abt re-fetching new OTF processings
        totbad = np.sum(sbimg[bad])

        # If all bad pixels have been zeroed, interpolate them.
        if totbad == 0:
            if verbose:
                print("WLIMAZ: {} bad pixels repaired".format(nbad))
            for i in range(nbad):
                xbad = bad[0][i] % ns
                ybad = bad[0][i] / ns
                if xbad != 0 and xbad != ns-1 and ybad != 0 and ybad != ns-1:
                    bad_val = (sbimg[ybad,xbad-1] + sbimg[ybad,xbad+1] +
                               sbimg[ybad-1,xbad] + sbimg[ybad+1,sbad])/4
                    sbimg[ybad,xbad] = bad_val

        indmx = np.unravel_index(np.argmax(sbimg, axis=None), sbimg.shape)
        xmx, ymx = indmx[1], indmx[0]
        xpos = xmx % ns
        ypos = ymx / ns

        # Threshold of 10 counts, FWHM of 2. Only return brightest result.
        star_finder = DAOStarFinder(10., 2., brightest=1)
        star_table = star_finder.find_stars(sbimg)
        star_x = star_table['xcentroid'][0]
        star_y = star_table['ycentroid'][0]
        if star_x < 0 or star_y < 0:
            print("WLIMAZ: Peak too close to edge: using approximate position.")
            star_x = xpos
            star_y = ypos

        star_x = star_x + xapprox - ns//2
        star_y = star_y + yapprox - ns//2
        if verbose:
            print("\tDAOFind: xc={}, yc={}".format(star_x, star_y))

        if nbad > 0:
            print("WLIMAZ: nbad={}, total bad counts={}".format(nbad, totbad))
        return star_x, star_y

def wlmeas(input_table, **kwargs):
    """
    Measure planetary nebula emission line locations.

    There are six planetary nebula emission lines that fall neatly into the WFC3 grism
    spectral orders, and this function uses the approximate wavelength solution to find
    the rough location of these lines, and then uses flux-weighting to determine the line
    centre. The user is able to override a given fit if the script is run with the
    show_plots flag.

    Parameters
    ----------
    input_table : abscal.common.exposure_data_table.AbscalDataTable
        Table of exposures with emission lines to fit.
    kwargs : dict
        Dictionary of overrides to the default reduction parameters, and command-line
        option selections.

    Returns
    -------
    output_table : astropy.table.Table
        Table of emission line locations
    """
    task = "wfc3: grism: wlmeas"
    default_values = get_defaults('abscal.common.args')
    base_defaults = default_values | get_defaults(kwargs.get('module_name', __name__))
    verbose = kwargs.get('verbose', base_defaults['verbose'])
    show_plots = kwargs.get('plots', base_defaults['plots'])
    if 'out_file' in kwargs:
        out_file = kwargs['out_file']
        out_dir, out_table = os.path.split(out_file)
        if out_dir == '':
            out_dir = os.getcwd()
    elif 'out_dir' in kwargs:
        out_dir = kwargs['out_dir']
    else:
        out_dir = os.getcwd()
    spec_name = kwargs.get('spec_dir', base_defaults['spec_dir'])
    spec_dir = os.path.join(out_dir, spec_name)
    in_notebook = kwargs.get('notebook', False)

    issues = {}
    exposure_parameter_file = get_data_file("abscal.wfc3", os.path.basename(__file__))
    if exposure_parameter_file is not None:
        with open(exposure_parameter_file, 'r') as inf:
            issues = yaml.safe_load(inf)

    fwhm = {
                'G102': {
                            -1: 80.,
                            1: 80.,
                            2: 40.
                        },
                'G141': {
                            -1: 120.,
                            1: 120.,
                            2: 60
                        }
            }

    # lab reference line wavelengths
    wl_air = np.array([9068.6, 9532.5, 10830., 12818.1, 16109.3, 16407.2])
    wl_vac = air2vac(wl_air)

    # Get wavelength reference file
    cal_data = get_data_file("abscal.wfc3", "calibration_files.yaml")
    with open(cal_data, 'r') as inf:
        cal_files = yaml.safe_load(inf)
    pn_ref_name = cal_files["ic5117_data"]
    pn_ref = get_data_file("abscal.wfc3", pn_ref_name)

    # Columns are 'wavelength' and 'flux'
    pn_ref_table = Table.read(pn_ref, format="ascii.basic")

    input_table = deepcopy(input_table)
    input_table = input_table[input_table["planetary_nebula"] == True]

    xpx = np.arange(1014, dtype=np.float64)

    if verbose:
        msg = "{}: Starting WFC3 wavelength measurement for GRISM data."
        print(msg.format(task))

    roots = []
    stars = []
    gratings = []
    x_ords = []
    y_ords = []
    orders = []
    lines = []
    notes = []

    for row in input_table:
        root = row["root"]
        path = row["path"]
        root_filter = input_table["root"]==row["root"]
        # Imaging files shouldn't be processed
        if row["filter"][0] == 'F':
            continue
        # Files without spectra (i.e. have no "extracted" entry) shouldn't be processed.
        if hasattr(row["extracted"], "mask") and row["extracted"].mask:
            continue

        spec_file = os.path.join(row['path'], row['extracted'])
        if not os.path.isfile(spec_file):
            # look for default extracted file
            extracted_name = "{}_{}_x1d.fits".format(row['root'], row['target'])
            print("Extracted Name: {}".format(extracted_name))
            extracted_dest = os.path.join(spec_dir, extracted_name)
            print("Extracted File: {}".format(extracted_dest))

            if os.path.isfile(extracted_dest):
                spec_file = extracted_dest
                extracted_value = os.path.join(spec_name, extracted_name)
                input_table['extracted'][root_filter] = extracted_value
                row['extracted'] = extracted_value
            else:
                msg = "{}: Unable to find extracted spectrum '{}'. Extracting."
                print(msg.format(task, row['extracted']))
                extract_table = input_table[input_table['root']==row['root']]
                output_row = reduce(extract_table, **kwargs)
                for item in ['path', 'extracted', 'xc', 'yc', 'xerr', 'yerr']:
                    input_table[item][root_filter] = output_row[item][0]
                    row[item] = output_row[item][0]
                spec_file = os.path.join(output_row['path'][0],
                                         output_row['extracted'][0])
                if not os.path.isfile(spec_file):
                    msg = "{}: {}: ERROR: EXTRACTION FAILED. SKIPPING ROW"
                    print(msg.format(task, row['root']))
                    continue
        # END search for the 1d extracted spectrum.

        if verbose:
            print("{}: reducing {}".format(task, root))
            print("\tpath is {}".format(path))
            print("\textracted is {}".format(row["extracted"]))
            print("\tfile is {}".format(spec_file))
            print("\trow: ")
            print(row)
        star = row["target"]
        grat = row["filter"]
        preamble = "wlmeas: {} ({}) ({})".format(root, grat, star)
        with fits.open(spec_file) as inf:
            wl = inf[1].data['wavelength']
            net = inf[1].data['net']
            dq = inf[1].data['eps'].astype(np.uint32)
            y_fit = inf[1].data['y_fit']
            if "xzorder" in inf[0].header:
                zxpos = inf[0].header['xzorder']
            else:
                zxpos = inf[0].header['xactual']
            if "yzorder" in inf[0].header:
                zypos = inf[0].header['yzorder']
            else:
                zypos = inf[0].header['yactual']

        # Bad DQ values are 256 (saturated), 516 (FF glitches),
        # and 32 (CTE tail) as per Table 2.5
        bad_dq = np.where((dq & (32|256|512)) != 0)
        good_dq = np.where((dq & (32|256|512)) == 0)

        name = row['root'][5:9]

        if grat == 'G102':
            wrang = [8900., 11100.]
        elif grat == 'G141':
            wrang = [10800., 17500.]
        else:
            raise ValueError("Unknown Grating/Filter {}".format(grat))

        pnvel = 0. # radial velocity
        target_star = find_star_by_name(row['target'])
        if target_star is not None:
            pnvel = target_star['radial_velocity']

        wrud = pn_ref_table['wavelength'] * (1. + pnvel/consts.c.to('km/s').value)
        wlref = wl_vac * (1. + pnvel/consts.c.to('km/s').value)

        for iord in [-1, 1, 2]:

            dlam = fwhm[grat][iord]

            # 3rd order 10830 at 32490 dominates 2nd order beyond ~16245
            # (32490/2)
            if iord == 2 and grat == 'G141':
                wrang[1] = 13000.

            xline = wlref * 0.

            line_dict = {}
            note_dict = {}

            for ilin in range(len(wlref)):
                wv = wl_vac[ilin]
                wl_line = wl/iord
                wlr = wlref[ilin]

                # All of these indicate that the line we're looking for isn't
                # in the order we're looking at.
                if wrang[1] > max(wl_line) or wv < wrang[0] or wv > wrang[1]:
                    continue

                # smoothed lines
                smorud = smooth_model(wrud, pn_ref_table['flux'], dlam)
                if grat == 'G102':
                    dw = 150/abs(iord)
                else:
                    dw = 250/abs(iord)

                xgdrud = np.where((wrud > wlr - dw) & (wrud < wlr + dw))

                xgood = np.where((wl/iord > wlr - dw) & (wl/iord < wlr + dw))

                min_x = wl[np.min(xgood)] - 50
                net_min = np.min(net[xgood])
                max_x = wl[np.max(xgood)] + 50
                net_max = np.min(net[xgood])

                if show_plots:
                    wl_good = np.where((wl > min_x-1000.) & (wl < max_x+1000.))

                    fig = plt.figure()
                    ax = fig.add_subplot(111)
                    title_str = "{} {} order={} line={} Search Range"
                    ax.set_title(title_str.format(grat, root, iord, wlr))
                    plt.plot(wl[wl_good], net[wl_good])
                    plt.plot([min_x, min_x], [net_min*0.5, net_max*1.5])
                    plt.plot([max_x, max_x], [net_min*0.5, net_max*1.5])
                    plt.show()

                fit_status = "auto"
                cmd = "unknown"
                maxpos = np.argmax(net[xgood])
                bad = np.where(dq[maxpos-1:maxpos+2] & (32|256|512) != 0)
                if len(bad) > 0:
                    nbad = len(bad[0])

                # This case handles the very bright 10830A line falling in
                #   the first order
                if ilin == 2 and iord == 1:
                    xcentr, ycentr = wlimaz(root, y_fit, wl, path, verbose)
                    xline[2] = xcentr
                    fit_status = "ima"

                if nbad > 0 and xline[ilin] == 0.:
                    print("WARNING: Centred using bad DQ")

                if root in issues:
                    siord = str(iord)
                    silin = str(ilin)
                    if siord in issues[root]:
                        if silin in issues[root][siord]:
                            xline[ilin] = float(issues[root][siord][silin]["xcentr"])
                            fit_status = "hardcoded"

                # Only do the centroiding if the line hasn't been set
                #   to a value yet.
                if xline[ilin] == 0 or fit_status == "ima":
                    xmin = np.min(xgood)
                    xmax = np.max(xgood)
                    search_range = (np.max(xgood) - np.min(xgood))/5
                    cont_low = np.where((xpx >= xmin) & (xpx <= xmin + search_range))
                    cont_hi = np.where((xpx >= xmax - search_range) & (xpx <= xmax))
                    cont = (np.mean(net[cont_low]) + np.mean(net[cont_hi]))/2
                    xcentr, fit = linecen(xgood[0], net[xgood], cont)
                    if fit == "good" or xline[ilin] == 0:
                        xline[ilin] = xcentr
                        if fit_status == "ima":
                            fit_status = "good (ima)"
                        else:
                            fit_status = fit

                cmd = {"choice": "unknown", "fit_status": fit_status, "msg": "",
                       "finished": False, "submitted": False}

                if show_plots or (("good" not in cmd["fit_status"]) and \
                                  ("hardcoded" not in cmd["fit_status"])):

                    while not cmd["finished"]:

                        cmd["submitted"] = False

                        fig, ax = plt.subplots()
                        fig.subplots_adjust(bottom=0.2)

                        title_str = "{} {} order={} line={}"
                        ax.set_title(title_str.format(grat, root, iord, wlr))
                        xrud = tabinv((wl/iord), wrud[xgdrud])
                        min_y = 0.
                        max_y = np.max(net[xgood])*1.1
#                         plt.plot(xrud, smorud[xgdrud], linestyle='--', label='WFC3 FWHM')
                        plt.scatter(wl[good_dq], net[good_dq], label='Good DQ')
                        plt.scatter(wl[bad_dq], net[bad_dq], label='Bad DQ')
                        if cmd["fit_status"] in ["custom", "rejected"]:
                            x_int = np.floor(cmd["choice"]).astype(np.int32)
                            xwl = wl[x_int] + (cmd["choice"]-x_int)*(wl[x_int+1]-wl[x_int])
                        else:
                            x_int = np.floor(xcentr).astype(np.int32)
                            xwl = wl[x_int] + (xcentr-x_int)*(wl[x_int+1]-wl[x_int])
                        if cmd["fit_status"] == "good" or cmd["fit_status"] == "good (ima)":
                            plt.plot([xwl, xwl], [0., 1.e10], color='green', label='Fit')
                        elif cmd["fit_status"] == "bad" or cmd["fit_status"] == "rejected":
                            plt.plot([xwl, xwl], [0., 1.e10], color='red',
                                     label='Bad Fit, Rejected Fit, or Not Found')
                        elif cmd["fit_status"] == "ima":
                            plt.plot([xwl, xwl], [0., 1.e10], color='blue',
                                     label='IMA zeroth read.')
                        elif cmd["fit_status"] == "hardcoded":
                            plt.plot([xwl, xwl], [0., 1.e10], color='grey',
                                     label='Hardcoded from Known Issues.')
                        elif cmd["fit_status"] == "custom":
                            plt.plot([xwl, xwl], [0., 1.e10], color='grey',
                                     label='User Custom')
                        else:
                            plt.plot([xwl, xwl], [0., 1.e10], color='red',
                                     label='UNKNOWN FIT')
                        plt.xlim(min_x, max_x)
                        plt.ylim(min_y, max_y)
                        plt.legend()

                        text_axes = fig.add_axes([0.75, 0.05, 0.15, 0.075])
                        msg = "Empty to accept fit, X to reject all fits, Wavelength "
                        msg += "for custom fit:"
                        text_box = TextBox(text_axes, msg, initial=cmd["msg"])

                        def submit(choice):
                            if choice == "":
                                cmd["finished"] = True
                                cmd["submitted"] = True
                            elif choice in ['x', 'X']:
                                wl_choice = wl[xgood][(len(xgood[0])-1)//2]
                                pix = np.searchsorted(wl, wl_choice)
                                cmd["choice"] = pix
                                cmd["fit_status"] = "rejected"
                                cmd["submitted"] = True
                            else:
                                try:
                                    custom_wl = float(choice)
                                    pix = np.searchsorted(wl, custom_wl)
                                    remainder = (custom_wl - wl[pix])/(wl[pix+1]-wl[pix])
                                    cmd["choice"] = pix + remainder
                                    cmd["fit_status"] = "custom"
                                    cmd["msg"] = ""
                                    cmd["submitted"] = True
                                except Exception as e:
                                    msg = "Enter nothing to accept fit, a floating point "
                                    msg += "value to choose a custom fit, or 'X' to reject"
                                    msg += "any fit."
                                    cmd["msg"] = msg
                                    cmd["fit_status"] = fit_status
                                    cmd["submitted"] = True
                            plt.close("all")

                        text_box.on_submit(submit)

                        def handle_close(evt):
                            if not cmd["submitted"]:
                                submit(text_box.text)

                        fig.canvas.mpl_connect('close_event', handle_close)
                        if in_notebook:
                            cmd["finished"] = True
                        plt.show()
                # end if showing the plot.

                if cmd["fit_status"] == "custom":
                    fit_status = "custom"
                    xcentr = cmd["choice"]
                elif cmd["fit_status"] == "rejected":
                    fit_status = "rejected"
                    xcentr = cmd["choice"]

                xline[ilin] = xcentr

                line_dict[int(wv)] = xcentr
                note_dict[int(wv)] = fit_status

                if verbose:
                    print("{}: Finished line {}".format(preamble, wv))

            roots.append(root)
            stars.append(star)
            gratings.append(grat)
            x_ords.append(zxpos)
            y_ords.append(zypos)
            orders.append(iord)
            lines.append(line_dict)
            notes.append(note_dict)

            if verbose:
                print("{}: Finished order {}.".format(preamble, iord))
        if verbose:
            print("{}: Finished obs {}".format(preamble, root))
    if verbose:
        print("{}: finished wavelength measurement.".format(preamble))

    output_table = Table()
    output_table['root'] = Column(data=roots)
    output_table['star'] = Column(data=stars)
    output_table['grism'] = Column(data=gratings)
    output_table['X_0_ORD'] = Column(data=x_ords)
    output_table['Y_0_ORD'] = Column(data=y_ords)
    output_table['order'] = Column(data=orders)
    for float_key in wl_vac:
        key = int(float_key)
        key_data = []
        for line_dict in lines:
            if key in line_dict:
                key_data.append(line_dict[key])
            else:
                key_data.append(-1.)
        note_data = []
        for note_dict in notes:
            if key in note_dict:
                note_data.append(note_dict[key])
            else:
                note_data.append("")
        output_table["{}_pos".format(key)] = Column(data=key_data, format='.3f')
        output_table["{}_notes".format(key)] = Column(data=note_data)

    return output_table

def wlmake(input_table, wl_table, **kwargs):
    """
    Derives a grism wavelength fit.

    Once planetary nebula emission lines have been located and fit, it is possible to use
    them as input in creating a full wavelength fit for the grism detector. ABSCAL fits
    the wavelength set with a linear slope and intercept, where both the slope and the
    intercept have constant terms, linear terms in X, and linear terms in y.

    Once the fit has been calculated, the script prints out fit errors based on the input
    exposures, and creates an output table with all of the fit terms. A separate fit is
    derived for each order of each grism.

    Parameters
    ----------
    input_table : abscal.common.exposure_data_table.AbscalDataTable
        Table of exposures to be fit.
    wl_table : astropy.table.Table
        Output of wlmeas.
    kwargs : dict
        Dictionary of overrides to the default reduction parameters, and command-line
        option selections.

    Returns
    -------
    output_table : astropy.table.Table
        Table of wavelength fit values
    """
    task = "wfc3: grism: wlmake"
    default_values = get_defaults('abscal.common.args')
    base_defaults = default_values | get_defaults(kwargs.get('module_name', __name__))
    verbose = kwargs.get('verbose', base_defaults['verbose'])
    show_plots = kwargs.get('plots', base_defaults['plots'])
    if 'out_file' in kwargs:
        out_file = kwargs['out_file']
        out_dir, out_table = os.path.split(out_file)
        if out_dir == '':
            out_dir = os.getcwd()
    elif 'out_dir' in kwargs:
        out_dir = kwargs['out_dir']
    else:
        out_dir = os.getcwd()
    spec_name = kwargs.get('spec_dir', base_defaults['spec_dir'])
    spec_dir = os.path.join(out_dir, spec_name)

    if verbose:
        print("Starting {}\nInput data:\n{}".format(task, wl_table))

    line_array = [9071, 9535, 10832, 12821, 16411] #16113, 16411]
    visible_lines = {
                        'G102': [9071, 9535, 10832],
                        'G141': [10832, 12821, 16411] #16113, 16411]
                    }
    wl_vac = np.array([[ 9070.0,  9070.0,  9071.4,  9070.5],
                       [ 9535.1,  9535.1,  9535.1,  9535.2],
                       [10833.5, 10834.6, 10833.4, 10833.4],
                       [12820.8, 12820.8, 12821.6, 12821.0],
                       # WHERE IS 16113?
#                       [],
                       [16413.2, 16414.0, 16412.1, 16412.7]])

    bad = np.where(wl_table['X_0_ORD'] == 0.)
    if len(bad) > 0:
        nbad = len(bad[0])

    results = {
                'grism': [],
                'order': [],
                'b_constant': [],
                'b_x': [],
                'b_y': [],
                'm_constant': [],
                'm_x': [],
                'm_y': []
              }

    for grism_index,grism in enumerate(['G102', 'G141']):
        if verbose:
            print("{}: starting grism {}".format(task, grism))

        mask = [g == grism for g in wl_table['grism']]
        grism_table = wl_table[mask]

        for iord in [-1, 1, 2]:
            if verbose:
                print("{}: {}: starting order {}".format(task, grism, iord))
            ord_mask = [o == iord for o in grism_table['order']]
            current_table = grism_table[ord_mask]
            xord_mask = [x > 0 for x in current_table['X_0_ORD']]
            current_table = current_table[xord_mask]
            emline_mask = []
            for row in current_table:
                is_any_valid = False
                for col in ["{}_pos".format(l) for l in line_array]:
                    if row[col] > 0:
                        is_any_valid = True
                emline_mask.append(is_any_valid)
            current_table = current_table[emline_mask]
            ngood = len(current_table)
            if ngood <= 0:
                continue
            if verbose:
                print("{}: {}: {}: Found {} good rows".format(task, grism, iord, ngood))
            dofil = current_table['root']
            dostr = current_table['star']
            doord = current_table['order'].data
            dozx = current_table['X_0_ORD'].data
            dozy = current_table['Y_0_ORD'].data
            wl_index = grism_index + 2*(abs(iord) - 1)
            results['grism'].append(grism)
            results['order'].append(iord)

            # Set up line offsets due to standard star radial velocity
            radial_velocity = np.zeros((ngood,), dtype=np.float64)
            rvc_p1 = np.zeros((ngood,), dtype=np.float64)
            for i,row in enumerate(current_table):
                star = find_star_by_name(row['star'])
                radial_velocity[i] = star['radial_velocity']
                rvc_p1[i] = 1 + star['radial_velocity']/consts.c.to('km/s').value

            good_fits = {}
            good_lines = []
            print(current_table)
            for line in visible_lines[grism]:
                print("Checking line {}".format(line),end='')
                good_fits[line] = 0
                for row in current_table:
                    notes_col = '{}_notes'.format(line)
                    if ('good' in row[notes_col]) or ('custom' in row[notes_col]):
                        good_fits[line] += 1
                        print(" good ",end='')
                    else:
                        print(" bad ",end='')
                print("{} good of {} ".format(good_fits[line], ngood),end='')
                if good_fits[line] > 0.5*ngood:
                    print("adding line {}.".format(line))
                    good_lines.append(line)
                else:
                    print("rejected.")
            nline = len(good_lines)
            low_line, high_line = min(good_lines), max(good_lines)
            low_idx, high_idx = line_array.index(low_line), line_array.index(high_line)
            doxi = np.zeros((ngood, nline), dtype=np.float64)
            for i,line in enumerate(good_lines):
                doxi[:,i] = current_table['{}_pos'.format(line)].data
            dox1 = current_table['{}_pos'.format(low_line)].data
            dox2 = current_table['{}_pos'.format(high_line)].data
            ref_line = np.full((ngood,), wl_vac[low_idx, wl_index])
            disp = (wl_vac[high_idx, wl_index] - wl_vac[low_idx, wl_index]) * doord
            disp *= rvc_p1/(dox2 - dox1)
            ref_line *= rvc_p1

            # Make polynomial fit to WL = b + m*delpx, spit out coefficient, and
            #   iterate to check results. Where delpx = x - x0, b = b1 + b2*x + b3*y,
            #   m = m1 + m2*x + m3*y, x0 is the z-order reference pixel location.
            xpx = np.arange(1014, dtype=np.float64)
            fit_good = np.where(((dox1 > 0.) & (dox2 > 0.)))
            n_fit_good = len(fit_good[0])

            #b3rd is 3rd element of b array. First 2 are X,Y of z-order.
            b3rd = ref_line[fit_good]*iord
            b3rd -= disp[fit_good]*(dox1[fit_good] - dozx[fit_good])
            b = np.zeros((n_fit_good,3), dtype=np.float64)
            b[:,0] = dozx[fit_good]
            b[:,1] = dozy[fit_good]
            m = deepcopy(b)
            b[:,2] = b3rd[:]
            if verbose:
                print("b matrix is {}".format(b))
            b_fit,_,_,_ = lstsq(np.c_[b[:,0], b[:,1], np.ones(b.shape[0])], b[:,2])
            b_x, b_y, b_const = b_fit
            bfit = [b_const + b_x*x + b_y*y for x,y in zip(b[:,0], b[:,1])]
            results['b_constant'].append(b_const)
            results['b_x'].append(b_x)
            results['b_y'].append(b_y)
            bval = b_const + b_x*506 + b_y*506 # b at (x,y) = (506,506)
            xr = [np.min(dozx[fit_good]), np.max(dozx[fit_good])] # range of 0-order points
            if verbose:
                print("{}: {}: {}: Z0 x-range is {}".format(task, grism, iord, xr))
                print("\t b1={}, b2={}, b3={}".format(b_const, b_x, b_y))
            m[:,2] = disp[fit_good]
            if verbose:
                print("m matrix is {}".format(m))
            m_fit,_,_,_ = lstsq(np.c_[m[:,0], m[:,1], np.ones(m.shape[0])], m[:,2])
            m_x, m_y, m_const = m_fit
            mfit = [m_const + m_x*x + m_y*y for x,y in zip(m[:,0], m[:,1])]
            results['m_constant'].append(m_const)
            results['m_x'].append(m_x)
            results['m_y'].append(m_y)
            mval = m_const + m_x*xpx * m_y*506
            mval = m_const + m_x*xpx * m_y*106
            mval = m_const + m_x*xpx * m_y*906
            if verbose:
                print("\t m1={}, m2={}, m3={}".format(m_const, m_x, m_y))

            line = []
            for em_line in visible_lines[grism]:
                line.append(wl_vac[line_array.index(em_line), wl_index])

            for ilin in range(nline):
                line_good = np.where((doord == iord) & (dox1 > 0.) & (dox2 > 0) & \
                                     (doxi[:,ilin] > 0))
                n_line_good = len(line_good[0])
                if n_line_good <= 0:
                    break # in theory break from the entire order, but we do what we can.
                wlerr = np.zeros((n_line_good,), np.float64)
                if verbose:
                    msg = "{}: {}: {}: File      Xmeas (px)  Xfit   err (A)"
                    print(msg.format(task, grism, iord))
                for igd in range(n_line_good):
                    indx = line_good[0][igd]
                    row = current_table[indx]
                    bval = b_const + b_x*dozx[indx] + b_y*dozy[indx]
                    mval = m_const + m_x*dozx[indx] + m_y*dozy[indx]
                    wnew = bval + mval*(xpx - dozx[indx])
                    xfit = tabinv(wnew, line[ilin]*iord*rvc_p1[indx])
                    wlerr[igd] = (doxi[indx,ilin] - xfit)*mval
                    if verbose:
                        msg = "                 {} {:8.2f} {:8.2f} {:8.2f} YZO={:8.2f}"
                        print(msg.format(row['root'], doxi[indx,ilin], xfit[0],
                                         wlerr[igd], dozy[indx]))
                if verbose:
                    msg = "Line, rms (A) and avg={}, {}, {}, #Obs={}"
                    print(msg.format(line[ilin], np.std(wlerr), np.mean(wlerr),
                                     n_line_good))

            dmeas = m[:,2]
            bmeas = b[:,2]
            xpos, ypos = dozx, dozy
            berr = bmeas - bfit
            merr = dmeas - mfit
            if verbose:
                msg = "{}: {}: {}: File     ZX (px)    ZY    b fit (A)     bmeas     berr"
                print(msg.format(task, grism, iord))
                for i in range(ngood):
                    row = current_table[i]
                    msg = "               {} {:8.2f} {:8.2f} {:8.2f} {:8.2f}  {:8.2f}"
                    print(msg.format(row['root'], xpos[i], ypos[i], bfit[i], bmeas[i],
                                     berr[i]))
                print("b rms={}".format(np.std(bmeas-bfit)))
                print("m rms={}".format(np.std(dmeas-mfit)))

            if show_plots:
                for igd in range(ngood):
                    row = current_table[igd]
                    full_row = input_table[input_table['root']==row['root']]
                    full_file = os.path.join(full_row["path"].data[0],
                                             full_row["extracted"].data[0])
                    if not os.path.isfile(full_file):
                        full_name = "{}_{}_x1d.fits".format(row['root'],
                                                            full_row['target'].data[0])
                        full_file = os.path.join(spec_dir, full_name)
                    if os.path.isfile(full_file):
                        star = find_star_by_name(row['star'])
                        rvc = star['radial_velocity']/consts.c.to('km/s').value
                        with fits.open(full_file) as inf:
                            wl = inf[1].data['wavelength']
                            net = inf[1].data['net']
                            angle = inf[0].header["angle"]

                        fig, ax = plt.subplots()
                        title_str = "{} {} order={} Fitted vs. Measured Line Positions"
                        ax.set_title(title_str.format(row['root'], grism, iord))

                        bval = b_const + b_x*dozx[igd] + b_y*dozy[igd]
                        mval = m_const + m_x*dozx[igd] + m_y*dozy[igd]
                        wnew = bval + mval*(xpx - dozx[igd])
                        plt.plot(wl/iord, net, 'b', label='Initial Wavelength Estimate')
                        plt.plot(wnew/iord, net, 'g', label='Fitted Wavelength')
                        for i in range(len(wl_vac[:,wl_index])):
                            plt.plot([wl_vac[i,wl_index], wl_vac[i,wl_index]*(1+rvc)],
                                     [0., 1.e5], 'r', linestyle='dashed')
                        msg = "Zero-order at ({:.2f},{:.2f})"
                        plt.figtext(0.2,0.2,msg.format(dozx[igd], dozy[igd]))
                        plt.legend()
                        # Figure out X and Y limits
                        if grism == "G102":
                            wave_low, wave_high = 9000, 11000
                        elif grism == "G141":
                            wave_low, wave_high = 10500, 17500
                        plt.xlim(wave_low, wave_high)
                        ind_low = np.searchsorted(wnew/iord, wave_low, side='left')
                        ind_high = np.searchsorted(wnew/iord, wave_high, side='right')
                        net_region = net[ind_low:ind_high]
                        plt.ylim(0, np.max(net_region)*1.1)
                        plt.show()

                        if verbose:
                            print("{}: {}: {}: {}".format(task, grism, iord, row['root']))
                            print("b={}, m={}".format(bval, mval))
                            print("Measured Dispersion = {}".format(dmeas[igd]))
                            print("Measured b={}".format(bmeas[igd]))
                            print("b,m errors={}, {}".format(bval-bmeas[igd], mval-dmeas[igd]))
                            print("Angle={}".format(angle))
            # done interactive plot
        # DONE ORDER LOOP
    # DONE GRISM LOOP

    output_table = Table()
    output_table['grism'] = Column(data=results['grism'])
    output_table['order'] = Column(data=results['order'])
    output_table['b_constant'] = Column(data=results['b_constant'])
    output_table['b_x'] = Column(data=results['b_x'])
    output_table['b_y'] = Column(data=results['b_y'])
    output_table['m_constant'] = Column(data=results['m_constant'])
    output_table['m_x'] = Column(data=results['m_x'])
    output_table['m_y'] = Column(data=results['m_y'])

    return output_table

def wl_offset(input_table, **kwargs):
    """
    Derives wavelength offsets for white dwarf exposures.

    Cross-correlates flux and net from white dwarf exposures against one another to derive
    offsets

    Parameters
    ----------
    input_table : abscal.common.exposure_data_table.AbscalDataTable
        Table of exposures to have offsets generated.
    kwargs : dict
        Dictionary of overrides to the default reduction parameters, and command-line
        option selections.

    Returns
    -------
    output_table : astropy.table.Table
        Updated table
    """
    task = "wfc3: grism: wloffset"
    default_values = get_defaults('abscal.common.args')
    base_defaults = default_values | get_defaults(kwargs.get('module_name', __name__))
    verbose = kwargs.get('verbose', base_defaults['verbose'])
    show_plots = kwargs.get('plots', base_defaults['plots'])
    if 'out_file' in kwargs:
        out_file = kwargs['out_file']
        out_dir, out_table = os.path.split(out_file)
        if out_dir == '':
            out_dir = os.getcwd()
    elif 'out_dir' in kwargs:
        out_dir = kwargs['out_dir']
    else:
        out_dir = os.getcwd()
    spec_name = kwargs.get('spec_dir', base_defaults['spec_dir'])
    spec_dir = os.path.join(out_dir, spec_name)

    return input_table

def additional_args(**kwargs):
    """
    Additional command-line arguments.

    Provides additional command-line arguments that are unique to the wavelength fitting
    process.

    Returns
    -------
    additional_args : dict
        Dictionary of tuples in the form (fixed,keyword) that can be passed to an argument
        parser to create a new command-line option
    """
    module_name = kwargs.get('module_name', __name__)
    base_defaults = get_defaults(module_name)

    additional_args = {}

    table_help = "The input metadata table to use."
    table_args = ['table']
    table_kwargs = {'help': table_help}
    additional_args['table'] = (table_args, table_kwargs)

    plots_help = "Include result plots while running."
    plots_args = ["-p", "--plots"]
    plots_kwargs = {'dest': 'plots', 'action': 'store_true',
                    'default': base_defaults['plots'], 'help': plots_help}
    additional_args['plots'] = (plots_args, plots_kwargs)

    return additional_args

def parse_args(**kwargs):
    """
    Parse command-line arguments.

    Gets the custom arguments from wavelength fitting, and passes them to the common
    command-line option function.

    Returns
    -------
    res : namespace
        parsed argument namespace
    """
    description_str = 'Process files from metadata table.'
    default_out_file = kwargs.get('default_input_file', 'dirirstare.log')
    default_in_file = kwargs.get('default_input_file', 'dirirstare.log')

    args = additional_args(**kwargs)

    res = parse(description_str, default_out_file, args, **kwargs)

    if res.paths is not None:
        if "," in res.paths:
            res.paths = res.paths.split(",")
        else:
            res.paths = [res.paths]
    else:
        res.paths = []

    if res.table is None:
        res.table = "dirtemp.log"

    if len(res.paths) == 0:
        res.paths.append(os.getcwd())

    return res

def main(do_measure=True, do_make=True, **kwargs):
    """
    Run the wavelength fitting function(s).

    Runs the wavelength fitting function(s) if called from the command line, with
    command-line arguments added in. Can run wlmeas, wlmake, or both.

    Parameters
    ----------
    do_measure : bool, default True
        Run wlmeas function
    do_make : bool, default, True
        Run wlmake function
    kwargs : dict
        Dictionary of parameters to override when running.
    """
    kwargs['default_output_file'] = 'wlmeastmp.log'
    parsed = parse_args(**kwargs)

    for key in kwargs:
        if hasattr(parsed, key):
            setattr(parsed, key, kwargs[key])

    input_table = AbscalDataTable(table=parsed.table,
                                  duplicates='both',
                                  search_str='',
                                  search_dirs=parsed.paths)

    measure_fname = parsed.out_file
    if do_measure:
        wl_calib_table = wlmeas(input_table, **vars(parsed), **kwargs)
        wl_calib_table.write(measure_fname, format='ascii.ipac', overwrite=True)
    else:
        wl_calib_table = Table.read(measure_fname, format='ascii.ipac')

    if do_make:
        final_wave_table = wlmake(input_table, wl_calib_table, **vars(parsed), **kwargs)
        (table_file, table_ext) = os.path.splitext(parsed.out_file)
        final_fname = table_file + "_final" + table_ext
        final_wave_table.write(final_fname, format='ascii.ipac', overwrite=True)
    # Done.

if __name__ == "__main__":
    main(module_name='abscal.wfc3.reduce_grism_wavelength')


import numpy as np
from .load_sed import *

class Simplemock():
    """
    Calculate flux through a filter for all stellar particle.

    Example
    >>> from general import defaults
    >>> import quick_mock
    >>> import galaxymodule
    >>> dfl = defaults.Default()
    >>> nout = 312#782 # 312
    >>> s = load.sim.Sim(nout=nout)
    >>> gcat = tree.halomodule.Halo(nout=nout, is_gal=True)
    >>> gg = galaxymodule.rd_GM.Gal(nout, catalog=gcat.data[21].copy(), info=s.info)
    >>> gg.debug=False
    >>> make_gal.mk_gal(gg)
    >>> from galaxymodule import quick_mock as qmc
    >>> MockSED = qmc.Simplemock(repo=dfl.dir_repo+'sed/')
    >>> gg.star.Flux_u= MockSED.get_flux(star=gg.star, filter_name='u')
    >>> gg.star.Flux_g= MockSED.get_flux(star=gg.star, filter_name='g')
    >>> gg.star.Flux_r= MockSED.get_flux(star=gg.star, filter_name='r')
    >>> gg.star.Flux_i= MockSED.get_flux(star=gg.star, filter_name='i')
    >>> gg.star.Flux_z= MockSED.get_flux(star=gg.star, filter_name='z')
    >>> Fluxs = [gg.star.Flux_u,
                 gg.star.Flux_g,
                 gg.star.Flux_r,
                 gg.star.Flux_i,
                 gg.star.Flux_z]
    >>> quick_mock.draw(Fluxs, gg.star["x"], gg.star["y"], suffix="face")
    >>> quick_mock.draw(Fluxs, gg.star["x"], gg.star["z"], suffix="edge")

    TODO

    t_univ

    interpolation outside range. (metallcity < 0.0004)

    """
    def __init__(self, repo="/home/hoseung/Work/pyclusterevol/repo/sed/",
                 filter_system="SDSS",
                 sed_model="bc03yy",
                 info=None,
                 load=True,
                 imf = "Salpeter"):
        self.filter_system = filter_system
        self.repo = repo
        self.sed_model = sed_model
        self.IMF = imf
        self.info = info
        if load is True:
            self.load_filters()
            self.load_SED_wavelength()
            self.load_SED_all()

    def load_filters(self):
        if self.filter_system == "SDSS":
            filter_lambda, filter_u, filter_g, filter_r, filter_i, filter_z = \
                np.genfromtxt(self.repo + "filter_sdss.dat",
                                skip_header=1, unpack=True)

            self.filters = {"lambda":filter_lambda,
                            "u":filter_u,
                            "g":filter_g,
                            "r":filter_r,
                            "i":filter_i,
                            "z":filter_z}

    def load_SED_wavelength(self):
        if self.sed_model == "bc03yy":
            self.sed_wavelength = np.genfromtxt(self.repo + "lambda.dat")

    def load_SED_all(self):
        """
        Full SEDs are just a few tens of MB.
        """
        if self.sed_model == "bc03yy":
            self.metal_points = np.array([0.0004, 0.001, 0.004, 0.01, 0.02, 0.04])
            # age points in tables.
            self.age_points = np.genfromtxt(self.repo+"ages_yybc.dat") # in Gry unit
            self.SEDs = np.zeros((6, 221, 1221))

            for i, metal in enumerate(self.metal_points):
                self.SEDs[i,:,:] = np.genfromtxt(self.repo +
                                            "bc03_yy_{:.4f}".format(metal)).reshape(221, 1221)
        elif self.sed_model == "bc":
            all_sed=[]
            for fn in glob("./BC03_v03/bc03/models/Padova2000/salpeter/*lr*.ised*.gz"):
                bc03 = Sed_block(fn)
                all_sed.append(bc03)
                print(bc03.isochrone)
                print(bc03.X, bc03.Y, bc03.Z)

            SED = Sed(all_sed)
            self.SEDs

        else:
            print("Sorry, Only bc03 is implemented.")

    def get_flux(self, star,
                 cell=None,
                 simple=False,
                 metal_lower_cut = True,
                 filter_name='r',
                 quick=False,
                 speed_check=False):
        """
        calculate SED of each particle.
        If cell is not None, attenuate flux accordingly.

        parameters
        ----------
        quick : False
            If True, reduce wavelength point of sed and filter significantly.
            This still gives reasonable value.

        speed_check : False
            If True, measure the time taken. Only for optimizing purpose.
        """
        if speed_check:
            from time import time
            t0 = time()

        Lum_sun = 3.826e33
        # BC2003 is in unit of L_sun Ang-1, where L_sun = Lum_sun.

        starmetal = star["metal"].copy() # Is the original array modified?
        starmetal[starmetal > 0.04] = 0.0399999 # a few stars have higher metallicity
        if metal_lower_cut:
            # No star with metallicity lower than the lowest table.
            starmetal[starmetal < min(self.metal_points)] = min(self.metal_points) * 1.0001

        locate_metal = np.digitize(starmetal, self.metal_points)-1 # GOOD
        relevant_metals = self.metal_points[:max(locate_metal)+2]
        nmetals = len(relevant_metals)

        # Star Age
        starage = star["age"] # the field name "time" should have been changed to "age"

        locate_age = np.digitize(starage, self.age_points)-1 # GOOD
        relevant_ages = self.age_points[:max(locate_age)+2]
        nages = len(relevant_ages)
        if speed_check: t1 = time() #

        ### Filter optimization. #################################################
        # Pick one
        this_filter = self.filters[filter_name]

        # band range
        i_filter_pos = this_filter > 0
        if quick:
            this_filter = this_filter[i_filter_pos][::10]
            filter_lambda_this_band = self.filters["lambda"][i_filter_pos][::10]
        else:
            this_filter = this_filter[i_filter_pos]
            filter_lambda_this_band = self.filters["lambda"][i_filter_pos]

        lambda_min_this_band = min(filter_lambda_this_band)
        lambda_max_this_band = max(filter_lambda_this_band)

        if quick:
            n_compress = 20
            sed_org = self.sed_wavelength
            sed_wavelength = self.sed_wavelength[::n_compress]
        else:
            sed_wavelength = self.sed_wavelength

        i_lambda_min = np.argmax(sed_wavelength > lambda_min_this_band) -1
        i_lambda_max = np.argmax(sed_wavelength > lambda_max_this_band)

        # Only a small part of SED is needed.
        # To compute d_lambda, one additional lambda point is desired.
        # Could be forward / backward / midpoint and so on.
        # let me take backward as fractional chnge in d_lambda is less in longer wavelength
        # Well.. actually I don't care..
        # d_lambda = wavelength[:-1] - wavelength[1:]
        wavelength = sed_wavelength[i_lambda_min:i_lambda_max+2] # why +2?
        n_wavelength = len(wavelength)-1#i_lambda_max - i_lambda_min + 1

        ##### Caclulate band flux #################
        # Load only necessary data
        # Load all once, keep under the class and copy a part of it when needed here.
        seds = np.zeros((nmetals, nages, n_wavelength)) # metal age lambda
        if self.sed_model == "bc03yy":
            for i, metal in enumerate(relevant_metals):
                if quick:
                    for j in range(seds.shape[1]):
                        seds[i,j,:] = np.interp(wavelength[:-1],
                                            sed_org,
                                            self.SEDs[i,j,:])
                else:
                    seds[i,:,:] = self.SEDs[i,:nages, i_lambda_min:i_lambda_max+1]

        if speed_check: t2 = time() # all set up

        # All are array calculations.
        # interpolation weight
        dl_m = (starmetal - relevant_metals[locate_metal] ) / \
                                     (relevant_metals[locate_metal+1] - relevant_metals[locate_metal])
        dr_m = (relevant_metals[locate_metal+1] - starmetal) / \
                                     (relevant_metals[locate_metal+1] - relevant_metals[locate_metal])
        dl_a = (starage - relevant_ages[locate_age] )   / \
                                     (relevant_ages[locate_age+1] - relevant_ages[locate_age])
        dr_a = (relevant_ages[locate_age+1] - starage ) / \
                                     (relevant_ages[locate_age+1] - relevant_ages[locate_age])

        if speed_check: t3 = time() # done first easy calculation

        # 2D linear interpolation
        # weight * SED.
        Flux =  np.multiply( (dr_m * dr_a), seds[locate_metal, locate_age,:].T).T +\
                np.multiply( (dl_m * dr_a), seds[locate_metal + 1, locate_age,:].T).T +\
                np.multiply( (dr_m * dl_a), seds[locate_metal, locate_age + 1, :].T).T +\
                np.multiply( (dl_m * dl_a), seds[locate_metal + 1, locate_age + 1,:].T).T

        if speed_check: t4 = time()
        # Convolve filter
        # Wavelengths at which filter function are defined are different from the SED wavelength points.
        # Interpolate filter function on SED points.
        filter_in_sed_wavelengths = np.interp(wavelength, filter_lambda_this_band, this_filter)
        Flux = np.multiply(filter_in_sed_wavelengths[:-1] * wavelength[-1], Flux)#\
        div = np.multiply(filter_in_sed_wavelengths[:-1], wavelength[-1])

        if speed_check: t5 = time()
        # Need to multiply stellar mass

        if cell is None or self.info is None or len(cell) == 0:
            return np.sum(Flux, axis=1) / np.sum(div) * Lum_sun * star["m"]
        else:
            #print(len(cell))
            colden = get_star_colden(star, cell) *self.info.unit_nH *self.info.unit_l / self.info.boxtokpc #* 1e3
            try:
                if colden == -1:
                    return np.sum(Flux, axis=1) / np.sum(div) * Lum_sun * star["m"]
            except:
                pass
            #lams = np.linspace(1e3,1e4,1e3) # in Angstrom
            waven = wavelength[:-1] * 1e-4 # in 1e-6 m

            Es = 0.44#*EBV
            # Hydrogen column number density
            colden = colden/5.8e21  # -1mag per 5.8e21 Hydrogen.
            tau = 0.4 * Es * np.outer(ext_curve_k(waven), colden)
            # colden = N_star array.
            # waven = N_wavelengths array
            # tau = N_wave X N_star
            # Flux = N_star X N_wave  ... (may be the other way around.)
            #print("tau", tau)
            F_ext= [ff*np.power(10,(-1*tt)) for ff, tt in zip(Flux, tau.T)]

            #print("before ext", np.sum(Flux, axis=1))
            #print("After ext", np.sum(F_ext, axis=1))
            if speed_check:
                t6 = time()
                print("age metal digitize {:.3f}".format(t1-t0))
                print("setup done {:.3f}".format(t2-t0))
                print("interpolation coefficient {:.3f}".format(t3-t0))
                print("flux {:.3f}".format(t4-t0))
                print("filter, flux, div {:.3f}".format(t5-t0))
                print("After dust attanuation {:.3f}".format(t6-t0))
            return np.sum(F_ext, axis=1) / np.sum(div) * Lum_sun * star["m"]

##################################################################
def flux2mag(flux,
             gal,
             x1="x",
             x2="y",
             filter_name = 'r',
             gal_range=None,
             Lum_dist = 400,
             plate_scale = 0.24,
             npixmax = 1200):
    # Observation conndition
    # in Mpc.
    if gal_range is None:
        gal_range = [[-gal.meta.rgal,gal.meta.rgal]]*2

    npixx, npixy = get_npix(plate_scale, gal_range, Lum_dist, npixmax)

    # Calculate Unit.
    d_lum_10p = 3.0857e19 # lumminonsity distance of 10pc in cm
    speed_of_light = 3e18 # angstrom / sec
    kpc_to_cm = 3.0857e21
    ldcm = Lum_dist * kpc_to_cm
    inv_distance = 1/(4*np.pi * ldcm * ldcm)

    band=BandSDSS()
    # Additional factors to derive realistic flux values.

    print(npixx, npixy)
    Flux_map = np.histogram2d(gal.star[x1], gal.star[x2],
               weights=flux,
               bins=[npixx,npixy],
               range=gal_range)[0]

    Flux_map *= Lum_sun * 1e-2 * inv_distance
    return - 2.5 * np.log10(Flux_map) \
           - 5. * np.log10(getattr(band, filter_name)["pivot_lambda"]) \
           + 2.5 * np.log10(speed_of_light) -48.6

def get_npix(plate_scale, gal_range, Lum_dist, npixmax):
    FOVx = (gal_range[0][1] - gal_range[0][0]) / (Lum_dist*1e3) * 180. / np.pi * 3600. # in arcsec
    FOVy = (gal_range[1][1] - gal_range[1][0]) / (Lum_dist*1e3) * 180. / np.pi * 3600.
    npixx= int(np.ceil(FOVx/plate_scale))
    npixy= int(np.ceil(FOVy/plate_scale))

    npixx = min([npixmax, npixx])
    npixy = min([npixmax, npixy])

    return (npixx, npixy)

class BandSDSS():
    def __init__(self):
        self.u = dict(pivot_lambda = 3557.0, name="u")
        self.g = dict(pivot_lambda = 4702.0, name="g")
        self.r = dict(pivot_lambda = 6175.0, name="r")
        self.i = dict(pivot_lambda = 7491.0, name="i")
        self.z = dict(pivot_lambda = 8946.0, name="z")

def composite_rgb(x,y, weight_r, weight_g, weight_b,
                  npix=100,
                  multiply_r=1.3,
                  multiply_g=1.1,
                  multiply_b=1.0,
                  log_scale = True,
                  range=None):
    import numpy as np
    from PIL import Image

    rgbArray = np.zeros((npix,npix,3))
    rgbArray[..., 0] = np.histogram2d(x, y,
               weights=weight_r, bins=npix, range=range)[0] * multiply_r
    rgbArray[..., 1] = np.histogram2d(x,y,
               weights=weight_g, bins=npix, range=range)[0] * multiply_g
    rgbArray[..., 2] = np.histogram2d(x,y,
               weights=weight_b, bins=npix, range=range)[0] * multiply_b

    if log_scale:
        rgbArray = np.log10(rgbArray+1)
    rgbArray = rgbArray/rgbArray.max() * 255
    #return rgbArray

    img = Image.fromarray(rgbArray.astype('uint8'))

    return img.rotate(90)
    #img.save('myimg.jpeg')

def draw(gal,
         x1="x",
         x2="y",
         suffix="edge",
         npix=200,
         gal_range=None,
         R="Flux_g",
         G="Flux_r",
         B="Flux_i",
         cr=1.0, cg=1.0, cb=3.0):
    """
    Parameers
    ---------
    gal:

    suffix:

    npix:
        default = 200
    gal_range:
        image 2d span in kpc.
    """

    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.colors import LogNorm

    channel_r = getattr(gal.star, R)
    channel_g = getattr(gal.star, G)
    channel_b = getattr(gal.star, B)

    if gal_range is None:
        gal_range = [[-gal.meta.rgal,gal.meta.rgal]]*2

    fig, axs = plt.subplots(2,3)
    fig.set_size_inches(8,6)

    titles=["u","g","r","i","z","composite"]
    for i, ax in enumerate(axs.ravel()[:5]):
        ax.hist2d(gal.star[x1], gal.star[x2],
               weights=np.log10(getattr(gal.star, "Flux_"+titles[i])+1),
               bins=npix,
               cmap=plt.cm.binary_r,
               norm=LogNorm(),
               range=gal_range)
        ax.set_aspect("equal")
        ax.set_title(titles[i])

    # Try scaling Flux_x arrays to make a better composite image
    comp_img = composite_rgb(gal.star[x1], gal.star[x2],
                             channel_r,
                             channel_g,
                             channel_b,
                             npix=npix,
                             multiply_r = cr,
                             multiply_g = cg,
                             multiply_b = cb,
                             log_scale=True,
                             range=gal_range)
    ax = axs[-1][-1]
    ax.imshow(comp_img)
    ax.set_title("composite")
    labels = [item.get_text() for item in ax.get_xticklabels()]
    empty_string_labels = ['']*len(labels)
    ax.set_aspect("equal")
    ax.set_xticklabels(empty_string_labels)
    ax.set_yticklabels(empty_string_labels)

    plt.savefig(str(gal.meta.id).zfill(5) + suffix + ".png", dpi=200)

def get_absolute_mag(flux, band=None, bandname="r"):
    """
    Don't forget to pass band.
    Initializing a new band class is expensive.

    Todo
    ----
    Figure out, why 1e-2??
    """
    if band is None:
        band = BandSDSS()

    speed_of_light = 3e18 # angstrom / sec
    d_lum_10p = 3.0857e19 # lumminonsity distance of 10pc in cm
    return - 2.5 * np.log10(np.sum(flux)/(4.*np.pi*d_lum_10p*d_lum_10p)*1e-2) \
            - 5. * np.log10(getattr(band, bandname)["pivot_lambda"]) \
           + 2.5 * np.log10(speed_of_light) - 48.6

def get_star_colden(star, cell):
    nstar = len(star)
    # no need to perform calculations on irrelevant cells
    # cut cells behind the further
    # This shrinks the cell length significantly.

    cell_dx_min = np.min(cell["dx"])

    ddx_h = cell["dx"] * 0.5

    sub_cell = cell[((cell["x"]+ddx_h) > star["x"].min()) * (cell["x"]-ddx_h < star["x"].max()) *\
                    ((cell["y"]+ddx_h) > star["y"].min()) * (cell["y"]-ddx_h < star["y"].max()) *\
                    ((cell["z"]+ddx_h) < star["z"].max())]

    #print("len subcell", len(sub_cell))
    if len(sub_cell) < 1:
        return -1

    ddx_h = sub_cell["dx"] * 0.5
    xl = sub_cell["x"] - ddx_h
    xr = sub_cell["x"] + ddx_h
    yl = sub_cell["y"] - ddx_h
    yr = sub_cell["y"] + ddx_h

    xrange = (xl.min(),xr.max())
    yrange = (yl.min(),yr.max())
    xspan = xrange[1]-xrange[0]
    yspan = yrange[1]-yrange[0]
    npixx = np.int(np.ceil(xspan/cell_dx_min))
    npixy = np.int(np.ceil(yspan/cell_dx_min))

    h = np.histogram2d(star["x"], star["y"],
                       bins=[npixx,npixy],
                       range=[xrange, yrange])

    dxmap = dymap = cell_dx_min

    # Sort cells
    i_relev = np.where(h[0].ravel() > 0)[0]

    # Sort stars
    ix=np.searchsorted(h[1], star["x"]) - 1
    iy=np.searchsorted(h[2], star["y"]) - 1

    ixy = iy + npixx*ix
    #print("histogram indicies and ixy are equivalent:", np.all(np.unique(ixy) == i_relev))

    i_star_sort = np.argsort(ixy)
    i_star = np.concatenate((np.searchsorted(ixy[i_star_sort], np.unique(ixy)),
                             [nstar]))

    # stars grouped in each lum map.
    sorted_star = star[i_star_sort]
    #sorted_flux = sorted_star["Flux_r"]#[i_star_sort]
    colden_star = np.zeros(nstar)

    #print('len star', len(i_star))
    #print("len cell", len(sub_cell))
    #print("minx cell", min(xl))
    #print("maxx cell", max(xr))
    #print("miny cell", min(yl))
    #print("maxy cell", max(yr))

    for i in range(len(i_star)-1):
        stars_here = sorted_star[i_star[i]:i_star[i+1]]
        jx, jy = ix[i_star_sort[i_star[i]]], iy[i_star_sort[i_star[i]]]
   #     print(jx, jy)
   #     print(h[1][jx], h[2][jy])
        cells_here = sub_cell[ (xr >= h[1][jx]) * (xl <= h[1][jx]) *\
                               (yr >= h[2][jy]) * (yl <= h[2][jy]) ]
   #     print(len(cells_here), "cells here")
        if len(cells_here) == 0:
            # at the edge
            colden_star[i_star[i]:i_star[i+1]] = 0
        else:
            # column density for each star.
            i_star_z = np.searchsorted(cells_here["z"]+cells_here["dx"], stars_here["z"])
            i_star_z[i_star_z >= len(cells_here)] = len(cells_here)-1
            colden_star[i_star[i]:i_star[i+1]] = np.cumsum(cells_here["var0"]*cells_here["dx"])[i_star_z]

    return colden_star

def ext_curve_k(lam, Rv=4.05):
    H_frac = 0.76

    lambda1 = 0.48613 # 1e-6 m
    lambda2 = 0.65628 # 1e-6 m

    # Why use only lambda 2
    inv_lambda1 = 1./lam[lam < lambda2]
    inv_lambda2 = 1./lam[lam > lambda2]

    k = np.zeros(len(lam))
    # lambda : 0.09 ~ 0.63 1e-6m
    k1 = 2.659 * (-2.156 + (1.509 * inv_lambda1)
                      - (0.198*inv_lambda1**2)
                  + (0.011*inv_lambda1**3)) + Rv
    # lambda : 0.63 ~ 5.08 1e-6m
    k2 = 2.659 * (-1.857 + (1.040*inv_lambda2)) + Rv

    k[lam < lambda2] = k1
    k[lam > lambda2] = k2
    return k