#! /opt/conda/bin/python3
""" Pipe class implementing the actual pipe simulation calculations """

# Copyright 2019 FAU-iPAT (http://ipat.uni-erlangen.de/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Optional
import numpy as np
from .base_pipe import BasePipe
from ..base.connectors.connector import Connector
from ..measure import Measure
from ..base.channels.import_channel import ImportChannel
from ..base.channels.export_channel import ExportChannel

class Pipe(BasePipe):  # pylint: disable=too-many-instance-attributes
    """
    Pipe class implementing the pipe simulation calculations
    """

    def __init__(
            self,
            diameter: float,
            length: float,
            wall_thickness: float,
            bulk_modulus: float,
            roughness: float,
            inner_points: int = None,
    ) -> None:
        """
        Initialization of the class

        :param diameter: Diameter of the fluid volume [m]
        :param length: Length of the pipe [m]
        :param wall_thickness: Thickness of the wall [m]
        :param bulk_modulus: Bulk modulus of wall material [Pa]
        :param roughness: Roughness of the wall [m]
        :param inner_points: Minimal number of inner points for discretization
        :raises TypeError: Wrong type of at least one parameter
        :raises ValueError: Value of at least one parameter out of bounds
        """
        super(Pipe, self).__init__(diameter, length, wall_thickness, bulk_modulus, roughness)
        if inner_points is not None and not isinstance(inner_points, int):
            raise TypeError('Wrong type for parameter inner_points ({} != {})'.format(type(inner_points), int))
        if inner_points is not None and inner_points < 3:
            raise ValueError('Number of inner points ({}) needs to greater than 2!'.format(inner_points))
        # Register internal fields
        self._inner_points = inner_points
        self._pressure: np.ndarray = self.field_create('pressure', 3)
        self._velocity: np.ndarray = self.field_create('velocity', 3)
        self.field_create('reynolds', 3)
        self.field_create('brunone', 3)
        self.field_create('darcy_friction_factor', 3)
        self._friction_steady = self.field_create('friction_steady', 3)
        self._friction_unsteady_a = self.field_create('friction_unsteady_a', 3)
        self._friction_unsteady_b = self.field_create('friction_unsteady_b', 3)
        self._sos: np.ndarray = self.field_create('speed_of_sound', 3)
        # Create the left connector
        self._left: Connector = Connector(self, [
            ExportChannel(Measure.deltaX, lambda: self._delta_x),
            ImportChannel(Measure.boundaryPoint, False),
            ExportChannel(Measure.diameter, lambda: self.diameter),
            ExportChannel(Measure.length, lambda: self.length),
            ExportChannel(Measure.area, lambda: self.area),
            ImportChannel(Measure.pressureLast, False),
            ExportChannel(Measure.pressureCurrent, lambda: self._pressure[0, 1]),
            ExportChannel(Measure.pressureLast, lambda: self._pressure[1, 1]),
            ImportChannel(Measure.velocityPlusCurrent, False),
            ImportChannel(Measure.velocityPlusLast, False),
            ExportChannel(Measure.velocityMinusCurrent, lambda: -self._velocity[0, 1]),
            ExportChannel(Measure.velocityMinusLast, lambda: -self._velocity[1, 1]),
            ExportChannel(Measure.frictionCurrent, lambda: self._friction_steady[0, 1] + self._friction_unsteady_b[0, 1]),
            ExportChannel(Measure.frictionLast, lambda: self._friction_steady[1, 1] + self._friction_unsteady_b[1, 1]),
            ExportChannel(Measure.BPspeedOfSoundCurrent, lambda: self._sos[0, 0]),
            ExportChannel(Measure.BPspeedOfSoundLast, lambda: self._sos[1, 0]),
        ])
        # Create the right connector
        self._right: Connector = Connector(self, [
            ExportChannel(Measure.deltaX, lambda: self._delta_x),
            ImportChannel(Measure.boundaryPoint, False),
            ExportChannel(Measure.diameter, lambda: self.diameter),
            ExportChannel(Measure.length, lambda: self.length),
            ExportChannel(Measure.area, lambda: self.area),
            ImportChannel(Measure.pressureLast, False),
            ExportChannel(Measure.pressureCurrent, lambda: self._pressure[0, -2]),
            ExportChannel(Measure.pressureLast, lambda: self._pressure[1, -2]),
            ImportChannel(Measure.velocityMinusCurrent, False),
            ImportChannel(Measure.velocityMinusLast, False),
            ExportChannel(Measure.velocityPlusCurrent, lambda: self._velocity[0, -2]),
            ExportChannel(Measure.velocityPlusLast, lambda: self._velocity[1, -2]),
            ExportChannel(Measure.frictionCurrent, lambda: self._friction_steady[0, -2] + self._friction_unsteady_a[0, -2]),
            ExportChannel(Measure.frictionLast, lambda: self._friction_steady[1, -2] + self._friction_unsteady_a[1, -2]),
            ExportChannel(Measure.BPspeedOfSoundCurrent, lambda: self._sos[0, -1]),
            ExportChannel(Measure.BPspeedOfSoundLast, lambda: self._sos[1, -1]),
        ])

    @property
    def left(self) -> Connector:
        """
        Left connector property

        :return: Left sided connector of the pipe
        """
        return self._left

    @property
    def right(self) -> Connector:
        """
        Right connector property

        :return: Right sided connector of the pipe
        """
        return self._right

    def get_max_delta_t(self) -> Optional[float]:
        """
        Method to return the maximum allowed timestep width for this component

        :return: Maximum allowed timestep width or None if any is suitable
        """
        n_min = self._inner_points if self._inner_points is not None else 3
        result = self.length / ((n_min + 1) * self.norm_speed_of_sound)
        return result

    def discretize(self, delta_t: float) -> None:
        """
        Method handling the discretization of the component (for a given timestep width)

        :param delta_t: Timestep width to discretize for
        :raises ValueError: Timestep too large to fit at least 3 inner points
        """
        self._delta_t = delta_t
        nodes = int(np.ceil(self.length / (self.norm_speed_of_sound * delta_t)) - 1)
        if nodes < 3:
            raise ValueError('Timestep to large!')
        self._delta_x = self.length / float(nodes + 1)
        self.fields_resize(nodes + 2)

    def initialize(self) -> None:
        """
        Initialize the internal state of the component (after discretization was called)
        """
        self.field('velocity')[:, :] = np.zeros(self.field('velocity').shape)[:, :]
        self.field('pressure')[:, :] = self.fluid.norm_pressure * np.ones(self.field('pressure').shape)[:, :]
        # Initialize derived properties
        for _ in range(2):
            self._calculate_reynolds()
            self._calculate_friction()
            self._calculate_speed_of_sound()
            self.fields_move()

    def prepare_next_timestep(self, delta_t: float, next_total_time: float) -> None:
        """
        Prepare the internal state for the next timestep to be calculated

        :param delta_t: Timestep width for the next timestep
        :param next_total_time: Total simulation time at the end of the next timestep
        """
        # Shift all internal fields
        self.fields_move()

    def exchange_last_boundaries(self) -> None:
        """
        Exchange the boundary values from the last time steps
        """
        # Exchange previous values with the left boundary
        self._pressure[1, 0] = self.left.value(Measure.pressureLast)
        self._velocity[1, 0] = self.left.value(Measure.velocityPlusLast)
        # Exchange previous values with the right boundary
        self._pressure[1, -1] = self.right.value(Measure.pressureLast)
        self._velocity[1, -1] = -self.right.value(Measure.velocityMinusLast)

    def finalize_current_timestep(self) -> None:
        """
        Method to perform final calculations at the end of the current timestep
        """
        # Exchange current values
        self._velocity[0, 0] = self.left.value(Measure.velocityPlusCurrent)
        self._velocity[0, -1] = -self.right.value(Measure.velocityMinusCurrent)
        # Calculate static values
        self._calculate_reynolds()
        self._calculate_friction()
        self._calculate_speed_of_sound()

    def prepare_next_inner_iteration(self, iteration: int) -> None:
        """
        Method to prepare the internal state for the next inner iteration of the current timestep

        :param iteration: Number of the next inner iteration to prepare for
        """

    def exchange_current_boundaries(self) -> None:
        """
        Exchange boundary values from the current time step
        """

    def calculate_next_inner_iteration(self, iteration: int) -> bool:
        """
        Method to do the calculations of the next inner iteration

        :param iteration: Number of the next inner iteration
        :return: Whether this component needs another inner iteration afterwards
        """
        self._calculate_pressure()
        self._calculate_velocity()
        return False

    def _calculate_speed_of_sound(self) -> None:
        """
        Calculate the current speed of sound
        """
        pressure = self.field_wide_slice('pressure', 0)
        result = self.speed_of_sound(pressure=pressure, temperature=None)
        self.field_wide_slice('speed_of_sound', 0)[:] = result[:]

    def _calculate_reynolds(self) -> None:
        """
        Calculate the Reynolds number based on the values from the previous time step
        """
        # Get the input fields
        pressure = self.field_wide_slice('pressure', 0)
        velocity = self.field_wide_slice('velocity', 0)
        # Calculate fluid properties
        viscosity = self.fluid.viscosity(temperature=None, shear_rate=None)
        density = self.fluid.density(pressure=pressure, temperature=None)
        # Calculate the reynolds number
        result = (density * np.abs(velocity) * self.diameter) / viscosity
        # Store/return the calculated result
        self.field_wide_slice('reynolds', 0)[:] = result[:]

    def _calculate_darcy_friction_factor(self) -> None:
        """
        Calculates darcy's friction coefficient within the pipe
        """
        # Get the input fields
        reynolds = self.field_wide_slice('reynolds', 0)
        result = np.ones(reynolds.shape)
        # Calculate the friction factor (low Re)
        selector = np.logical_and(reynolds > 0.0, reynolds < 2100.0)
        if np.sum(selector) > 0:
            local_reynolds = reynolds[selector]
            factor = 64.0 / local_reynolds
            result[selector] = factor
        # Calculate the friction factor (high Re)
        selector = (reynolds >= 2100.0)
        if np.sum(selector) > 0:
            local_reynolds = reynolds[selector]
            factor = 10.0 * np.ones(local_reynolds.shape)
            error = np.ones(local_reynolds.shape)
            while np.any(error > 1e-12):
                term1 = self.roughness / (3.7 * self.diameter)
                term2 = 2.51 / (local_reynolds * np.sqrt(factor))
                temp = -2.0 * np.log10(term1 + term2)
                old_factor, factor = factor, np.square(1.0 / temp)
                error = np.abs(factor - old_factor)
            result[selector] = factor
        # Store/return the calculated result
        self.field_wide_slice('darcy_friction_factor', 0)[:] = result[:]

    def _calculate_friction_steady(self) -> None:
        """
        Calculate the steady friction using darcy's factor
        """
        # Get the input fields
        velocity = self.field_wide_slice('velocity', 0)
        friction_factor = self.field_wide_slice('darcy_friction_factor', 0)
        # Calculate the friction
        result = (friction_factor / (2.0 * self.diameter)) * np.abs(velocity) * velocity
        # Store/return the calculated result
        self.field_wide_slice('friction_steady', 0)[:] = result[:]

    def _calculate_friction(self) -> None:
        """
        Calculate the total friction (steady + unsteady)
        """
        # Calculate steady friction
        self._calculate_darcy_friction_factor()
        self._calculate_friction_steady()
        # Calculate unsteady friction
        self._calculate_brunone()
        self._calculate_unsteady_friction_a()
        self._calculate_unsteady_friction_b()

    def _calculate_pressure(self) -> None:
        """
        Calculate the pressure of the current time step
        """
        # Get the input fields
        pressure_center = self.field_slice('pressure', 1, 0)
        pressure_a = self.field_slice('pressure', 1, -1)
        pressure_b = self.field_slice('pressure', 1, +1)
        velocity_a = self.field_slice('velocity', 1, -1)
        velocity_b = self.field_slice('velocity', 1, +1)
        friction_a = self.field_slice('friction_steady', 1, -1) + self.field_slice('friction_unsteady_a', 1, -1)
        friction_b = self.field_slice('friction_steady', 1, +1) + self.field_slice('friction_unsteady_b', 1, +1)
        # Calculate fluid properties
        speed_of_sound = self.speed_of_sound(pressure=pressure_center, temperature=None)
        density = self.fluid.density(pressure=pressure_center, temperature=None)
        # Calculate the reynolds number
        result = 0.5 * (
            (speed_of_sound * density * (velocity_a - velocity_b))
            + (pressure_a + pressure_b)
            + (self._delta_t * speed_of_sound * density * (friction_b - friction_a))
            # todo: height terms
        )
        # Store/return the calculated result
        self.field_slice('pressure', 0, 0)[:] = result[:]

    def _calculate_velocity(self) -> None:
        """
        Calculate the velocity of the current time step
        """
        # Get the input fields
        pressure_center = self.field_slice('pressure', 1, 0)
        pressure_a = self.field_slice('pressure', 1, -1)
        pressure_b = self.field_slice('pressure', 1, +1)
        velocity_a = self.field_slice('velocity', 1, -1)
        velocity_b = self.field_slice('velocity', 1, +1)
        friction_a = self.field_slice('friction_steady', 1, -1) + self.field_slice('friction_unsteady_a', 1, -1)
        friction_b = self.field_slice('friction_steady', 1, +1) + self.field_slice('friction_unsteady_b', 1, +1)
        # Calculate fluid properties
        speed_of_sound = self.speed_of_sound(pressure=pressure_center, temperature=None)
        density = self.fluid.density(pressure=pressure_center, temperature=None)
        # Calculate the reynolds number
        result = 0.5 * (
            (velocity_a + velocity_b)
            + ((1.0 / (speed_of_sound * density)) * (pressure_a - pressure_b))
            - (self._delta_t * (friction_a + friction_b))
            # todo: height terms
        )
        # Store/return the calculated result
        self.field_slice('velocity', 0, 0)[:] = result[:]

    def _calculate_brunone(self) -> None:
        """
        Calculate the Brunone factor for unsteady friction
        """
        # Get the input fields
        reynolds = self.field_wide_slice('reynolds', 0)
        # Calculate the Brunone factor
        result = 0.000476 * np.ones(reynolds.shape)
        selector = (reynolds >= 2320.0)
        if np.sum(selector) > 0:
            local_reynolds = reynolds[selector]
            factor = 14.3 / np.power(local_reynolds, 0.05)
            factor = 7.41 / np.power(local_reynolds, np.log10(factor))
            result[selector] = factor
        result = np.sqrt(result) / 2.0
        # Store/return the calculated result
        self.field_wide_slice('brunone', 0)[:] = result[:]

    def _calculate_unsteady_friction_a(self) -> None:
        """
        Calculate the unsteady friction to left side
        """
        # Get the input fields
        brunone = self.field_ext_slice('brunone', 0, 0)
        velocity_a = self.field_ext_slice('velocity', 0, 0)
        velocity_aa = self.field_ext_slice('velocity', 1, 0)
        velocity_p = self.field_ext_slice('velocity', 0, 1)
        pressure_a = self.field_ext_slice('pressure', 0, 0)
        # Calculate fluid properties
        speed_of_sound = self.speed_of_sound(pressure=pressure_a, temperature=None)
        # Calculate the friction
        vdt = (velocity_a - velocity_aa) / self._delta_t
        vdx = (velocity_p - velocity_a) / self._delta_x
        result = brunone * (vdt + (speed_of_sound * np.sign(velocity_a * vdx) * vdx))
        # Store/return the calculated result
        self.field_ext_slice('friction_unsteady_a', 0, 0)[:] = result[:]

    def _calculate_unsteady_friction_b(self) -> None:
        """
        Calculate the unsteady friction to right side
        """
        # Get the input fields
        brunone = self.field_ext_slice('brunone', 0, 1)
        velocity_b = self.field_ext_slice('velocity', 0, 1)
        velocity_bb = self.field_ext_slice('velocity', 1, 1)
        velocity_p = self.field_ext_slice('velocity', 0, 0)
        pressure_b = self.field_ext_slice('pressure', 0, 1)
        # Calculate fluid properties
        speed_of_sound = self.speed_of_sound(pressure=pressure_b, temperature=None)
        # Calculate the friction
        vdt = (velocity_b - velocity_bb) / self._delta_t
        vdx = (velocity_b - velocity_p) / self._delta_x
        result = brunone * (vdt + (speed_of_sound * np.sign(velocity_b * vdx) * vdx))
        # Store/return the calculated result
        self.field_ext_slice('friction_unsteady_b', 0, 1)[:] = result[:]


# vim: fdm=marker
'''
author:     Fabio Zanini/Richard Neher
date:       25/04/2015
content:    Data access module HIV patients.
'''
# Modules
import numpy as np
import pandas as pd
from Bio import SeqIO, AlignIO
from .sequence import alpha, alphaa
from .filenames import (get_custom_reference_filename,
                        get_subtype_reference_alignment_filename,
                        get_subtype_reference_allele_frequencies_filename)

class ReferenceTranslator(object):
    """docstring for ReferenceTranslater"""
    def __init__(self, ref1='HXB2', ref2='NL4-3'):
        super(ReferenceTranslator, self).__init__()
        self.ref1 = ref1
        self.ref2 = ref2

        self.refseq1 = SeqIO.read(get_custom_reference_filename(self.ref1, format='gb'), format='genbank').seq
        self.refseq2 = SeqIO.read(get_custom_reference_filename(self.ref2, format='gb'), format='genbank').seq

        from seqanpy import align_global
        (score, ali1, ali2) = align_global(str(self.refseq1), str(self.refseq2), band=200)
        self.count1 = np.cumsum(np.fromstring(ali1,'S1')!='-')-1
        self.count2 = np.cumsum(np.fromstring(ali2,'S1')!='-')-1

    def translate(self, pos, ref=None):
        if ref is None:
            ref = self.ref1
        if ref == self.ref1:
            ii = np.searchsorted(self.count1,pos)
            if ii<len(self.count2):
                return self.ref2, self.count2[ii]
            else:
                return self.ref2, -1
        elif ref ==self.ref2:
            ii = np.searchsorted(self.count2,pos)
            if ii<len(self.count1):
                return self.ref1, self.count1[ii]
            else:
                return self.ref1, -1
        else:
            print("unknown reference", ref)
            return "not found", np.nan

class HIVreference(object):
    """docstring for HIVreference"""
    def __init__(self, refname='HXB2', subtype='B', load_alignment=True):
        self.refname = refname
        self.subtype = subtype
        self.seq = SeqIO.read(get_custom_reference_filename(self.refname, format='gb'), format='genbank')
        # translate genbank encoded sequence features into a dictionary
        self.annotation = {x.qualifiers['note'][-1]:x for x in self.seq.features}

        if load_alignment:
            fn = get_subtype_reference_alignment_filename(subtype=subtype,
                                                          refname=refname)
            self.aln = np.array(AlignIO.read(fn, 'fasta'))
            self.calc_nucleotide_frequencies()

        else:
            fn = get_subtype_reference_allele_frequencies_filename(subtype=subtype,
                                                                   refname=refname)
            self.af = np.load(fn)

        self.consensus_indices = np.argmax(self.af, axis=0)
        self.consensus = alpha[self.consensus_indices]
        self.calc_entropy()

    def calc_nucleotide_frequencies(self):
        self.af = np.zeros((len(alpha)-1, self.aln.shape[1]), dtype=float)
        for ni, nuc in enumerate(alpha[:-1]):
            self.af[ni,:] = np.sum(self.aln==nuc.astype("U1"), axis=0)
        cov = np.sum(self.af, axis=0)
        self.af /= cov + 1e-10

    def calc_entropy(self):
        self.entropy = np.maximum(0,-np.sum(self.af*np.log(1e-10+self.af), axis=0))

    def map_to_sequence_collection():
        pass

    def get_ungapped(self, threshold=0.05):
        return self.af[-1,:] < 0.05

    def get_entropy_quantiles(self, q):
        from scipy.stats import scoreatpercentile
        thresholds = [scoreatpercentile(self.entropy, 100.0*i/q) for i in range(q+1)]
        return {i: {'range':(thresholds[i],thresholds[i+1]),
                    'ind':np.where((self.entropy>=thresholds[i])*(self.entropy<thresholds[i+1]))[0]}
               for i in range(q)}

    def get_entropy_in_patient_region(self, map_to_ref):
        '''
        returns entropy in a specific regions defined by a set of indices in the reference
        params:
        map_to_ref  --  either a one dimensional vector specifying indices in the reference
                        or a (3, len(region)) array with the reference coordinates in the first column
                        this is the output of Patient.map_to_external_reference
        '''
        if len(map_to_ref.shape) == 2:
            return self.entropy[map_to_ref[:, 0]]
        elif len(map_to_ref.shape) == 1:
            return self.entropy[map_to_ref]

    def get_consensus_in_patient_region(self, map_to_ref):
        '''
        returns consensus in a specific regions defined by a set of indices in the reference
        params:
        map_to_ref  --  either a one dimensional vector specifying indices in the reference
                        or a (3, len(region)) array with the reference coordinates in the first column
                        this is the output of Patient.map_to_external_reference
        '''
        if len(map_to_ref.shape) == 2:
            return self.consensus[map_to_ref[:, 0]]
        elif len(map_to_ref.shape) == 1:
            return self.consensus[map_to_ref]

    def get_consensus_indices_in_patient_region(self, map_to_ref):
        '''
        returns consensus_indices in a specific regions defined by a set of indices in the reference
        params:
        map_to_ref  --  either a one dimensional vector specifying indices in the reference
                        or a (3, len(region)) array with the reference coordinates in the first column
                        this is the output of Patient.map_to_external_reference
        '''
        if len(map_to_ref.shape) == 2:
            return self.consensus_indices[map_to_ref[:, 0]]
        elif len(map_to_ref.shape) == 1:
            return self.consensus_indices[map_to_ref]

class HIVreferenceAminoacid(object):
    def __init__(self, region, refname='HXB2', subtype='B'):
        self.region = region
        self.refname = refname
        self.subtype = subtype

        seq = SeqIO.read(get_custom_reference_filename(self.refname, format='gb'), format='genbank')
        # translate genbank encoded sequence features into a dictionary
        annotation = {x.qualifiers['note'][-1]:x for x in seq.features}
        self.seq = annotation[region].extract(seq)

        fn = get_subtype_reference_alignment_filename(region=region,
                                                      refname=refname,
                                                      subtype=self.subtype,
                                                      type='aa')
        self.aln = np.array(AlignIO.read(fn, 'fasta'))
        self.calc_aminoacid_frequencies()

        self.consensus_indices = np.argmax(self.af, axis=0)
        self.consensus = alphaa[self.consensus_indices]
        self.calc_entropy()

    def calc_aminoacid_frequencies(self):
        self.af = np.zeros((len(alphaa)-1, self.aln.shape[1]), dtype=float)
        for ai, aa in enumerate(alphaa[:-1]):
            self.af[ai,:] = np.sum(self.aln==aa, axis=0)
        cov = np.sum(self.af, axis=0)
        self.af /= cov

    def calc_entropy(self):
        self.entropy = np.maximum(0,-np.sum(self.af*np.log(1e-10+self.af), axis=0))

    def get_ungapped(self, threshold=0.05):
        return self.af[-1,:] < 0.05

    def get_entropy_quantiles(self, q):
        from scipy.stats import scoreatpercentile
        thresholds = [scoreatpercentile(self.entropy, 100.0*i/q) for i in range(q+1)]
        return {i: {'range':(thresholds[i],thresholds[i+1]),
                    'ind':np.where((self.entropy>=thresholds[i])*(self.entropy<thresholds[i+1]))[0]}
               for i in range(q)}

    def get_entropy_in_patient_region(self, map_to_ref):
        '''
        returns entropy in a specific regions defined by a set of indices in the reference
        params:
        map_to_ref  --  either a one dimensional vector specifying indices in the reference
                        or a (2, len(region)) array with the reference coordinates in the first column
                        this is the output of Patient.map_to_external_reference_aminoacids
        '''
        if len(map_to_ref.shape) == 2:
            return self.entropy[map_to_ref[:, 0]]
        elif len(map_to_ref.shape) == 1:
            return self.entropy[map_to_ref]

    def get_consensus_in_patient_region(self, map_to_ref):
        '''
        returns consensus in a specific regions defined by a set of indices in the reference
        params:
        map_to_ref  --  either a one dimensional vector specifying indices in the reference
                        or a (2, len(region)) array with the reference coordinates in the first column
                        this is the output of Patient.map_to_external_reference_aminoacids
        '''
        if len(map_to_ref.shape) == 2:
            return self.consensus[map_to_ref[:, 0]]
        elif len(map_to_ref.shape) == 1:
            return self.consensus[map_to_ref]

    def get_consensus_indices_in_patient_region(self, map_to_ref):
        '''
        returns consensus_indices in a specific regions defined by a set of indices in the reference
        params:
        map_to_ref  --  either a one dimensional vector specifying indices in the reference
                        or a (2, len(region)) array with the reference coordinates in the first column
                        this is the output of Patient.map_to_external_reference_aminoacids
        '''
        if len(map_to_ref.shape) == 2:
            return self.consensus_indices[map_to_ref[:, 0]]
        elif len(map_to_ref.shape) == 1:
            return self.consensus_indices[map_to_ref]


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2018 University of Groningen
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import defaultdict
import itertools
import networkx as nx

from .utils import maxes, first_alpha

def add_element_attr(molecule):
    for node_idx in molecule:
        node = molecule.nodes[node_idx]
        if 'element' not in node:
            try:
                element = first_alpha(node['atomname'])
            except KeyError:
                raise ValueError('Cannot guess the element of atom {}: '
                                 'the node has no atom name.'
                                 .format(node_idx))
            except ValueError:
                raise ValueError('Cannot guess the element of atom {}: '
                                 'the atom name has no alphabetic charater.'
                                 .format(node_idx))
            node['element'] = element

def categorical_cartesian_product(graph1, graph2, attributes=tuple()):
    product = nx.Graph()  # FIXME graphtype?
    for idx1, idx2 in itertools.product(graph1, graph2):
        node1 = graph1.nodes[idx1]
        node2 = graph2.nodes[idx2]
        if all(attr in node1 and attr in node2 and node1[attr] == node2[attr] for attr in attributes):
            attrs = {}
            for attr in set(node1.keys()) | set(node2.keys()):
                attrs[attr] = (node1.get(attr, None), node2.get(attr, None))
            product.add_node((idx1, idx2), **attrs)
    return product

def categorical_modular_product(graph1, graph2, attributes=tuple()):
    product = categorical_cartesian_product(graph1, graph2, attributes)
    for (graph1_node1, graph2_node1), (graph1_node2, graph2_node2) in itertools.combinations(product.nodes(), 2):
        graph1_nodes = graph1_node1, graph1_node2
        graph2_nodes = graph2_node1, graph2_node2
        both_edge = graph1.has_edge(*graph1_nodes) and graph2.has_edge(*graph2_nodes)
        neither_edge = not graph1.has_edge(*graph1_nodes) and\
                       not graph2.has_edge(*graph2_nodes)
        # Effectively: not (graph1.has_edge(graph1_node1, graph1_node2) xor graph2.has_edge(graph2_node1, graph2_node2))
        if graph1_node1 != graph1_node2 and graph2_node1 != graph2_node2 and\
                (both_edge or neither_edge):
            attrs = {}
            if both_edge:
                g1_edge_keys = set(graph1.edges[graph1_nodes].keys())
                g2_edge_keys = set(graph2.edges[graph2_nodes].keys())
                for attr in g1_edge_keys | g2_edge_keys:
                    attrs[attr] = (graph1.edges[graph1_nodes].get(attr, None),
                                   graph2.edges[graph2_nodes].get(attr, None))
            product.add_edge((graph1_node1, graph2_node1), (graph1_node2, graph2_node2), **attrs)
    return product

def categorical_maximum_common_subgraph(graph1, graph2, attributes=tuple()):
    product = categorical_modular_product(graph1, graph2, attributes)
    cliques = nx.find_cliques(product)
    # cliques is an iterator which will return a *lot* of items. So make sure
    # we never turn it into a full list.
    largest = maxes(cliques, key=len)
    matches = [dict(clique) for clique in largest]
    return matches

def maximum_common_subgraph(graph1, graph2, attributes=tuple()):
    product = nx.Graph()
    # First, find the MCS between all nodes of degree != 1, such as the carbons
    # Nothing new or exciting here.
    for g1_node, g2_node in itertools.product(graph1, graph2):
        node1 = graph1.nodes[g1_node]
        node2 = graph2.nodes[g2_node]
        if all(attr in node1 and attr in node2 and node1[attr] == node2[attr] for attr in attributes):
            if graph1.degree(g1_node) != 1 and graph2.degree(g2_node) != 1:
                product.add_node((g1_node, g2_node))
    for (g1_node1, g2_node1), (g1_node2, g2_node2) in itertools.combinations(product.nodes(), 2):
        both_edge = graph1.has_edge(g1_node1, g1_node2) and graph2.has_edge(g2_node1, g2_node2)
        neither_edge = not graph1.has_edge(g1_node1, g1_node2) and not graph2.has_edge(g2_node1, g2_node2)
        # Effectively: not (graph1.has_edge(g1_node1, g1_node2) xor graph2.has_edge(g2_node1, g2_node2))
        if g1_node1 != g1_node2 and g2_node1 != g2_node2 and (both_edge or neither_edge):
            product.add_edge((g1_node1, g2_node1), (g1_node2, g2_node2))
    cliques = nx.find_cliques(product)
#    largest = maxes(cliques, key=len)
    # We can't do maxes, because it might still grow to be as large. Does make
    # things slower though... We could say that it can grow to be at most the
    # current size plus the number of degree-1 nodes.
    largest = cliques

    # Add an empty match in case nothing of degree > 1 matches. In that case we
    # still need to do the loop below.
    largest = itertools.chain([[]], largest)

    # Now, for every MCS we found, look at the nodes of degree 1. The
    # attributes still need to match. In addition, they need to have the same
    # (mapped) neighbour, or the neighbour must be missing from the graph2 graph
    all_cliques = []
    for clique in largest:
        match = dict(clique)
        product = nx.Graph()
        product.add_nodes_from(clique)
        for g1_node, g2_node in itertools.product(graph1, graph2):
            node1 = graph1.nodes[g1_node]
            node2 = graph2.nodes[g2_node]
            # We can't do this above, because we need the match to translate
            # nodes from graph graph1 to graph graph2 to see whether their neighbours
            # correspond.
            if (graph1.degree(g1_node) <= 1 or graph2.degree(g2_node) <= 1) and\
                    all(attr in node1 and attr in node2 and node1[attr] == node2[attr] for attr in attributes):
                g1_neighbors = [match.get(n, None) for n in graph1.neighbors(g1_node)]
                # If no neighbors are found for g1_node, or if any of them are
                # the same in graph2, they're compatible.

                # FIXME?
                # This eliminates some nodes from the MCS, since it's possible
                # the MCS does not include *all* nodes in match. This means
                # that some nodes should be considered compatible, even if they
                # have different neighbours, but only if that neighbor is not
                # part of the final MCS. This makes testing a little tricky,
                # since categorical_maximum_common_subgraph *does* find them.
                # It'll be a cornercase anyway.
                if not g1_neighbors or None in g1_neighbors or any(n in g1_neighbors for n in graph2.neighbors(g2_node)):
                    product.add_node((g1_node, g2_node))
        for (g1_node1, g2_node1), (g1_node2, g2_node2) in itertools.combinations(product.nodes(), 2):
            both_edge = graph1.has_edge(g1_node1, g1_node2) and graph2.has_edge(g2_node1, g2_node2)
            neither_edge = not graph1.has_edge(g1_node1, g1_node2) and not graph2.has_edge(g2_node1, g2_node2)
            # Effectively: not (graph1.has_edge(g1_node1, g1_node2) xor graph2.has_edge(g2_node1, g2_node2))
            if g1_node1 != g1_node2 and g2_node1 != g2_node2 and (both_edge or neither_edge):
                product.add_edge((g1_node1, g2_node1), (g1_node2, g2_node2))
        # TODO: This duplicates a lot of effort. Maybe create the compatibility
        # graph first from all cliques, and find the cliques only once?
        this_pass = nx.find_cliques(product)
        all_cliques.append(this_pass)
    # And finally, find the largest MCS in all cliques graph2.
    largest = maxes(itertools.chain(*all_cliques), key=len)
    matches = set(frozenset(m) for m in largest)  # remove duplicates
    matches = [dict(clique) for clique in matches]

    return matches

def isomorphism(reference, residue):
    """
    Finds matching atoms between ``reference`` and ``residue``. ``residue`` should be
    a subgraph of ``reference``. Matchin is done based on connectivity and
    the ``element`` attribute of the nodes.

    The subgraph isomorphism is first calculated using non-hydrogen atoms only.
    These matches are then extended to include one option for the hydrogen
    isomorphism. This is done because otherwise a combinatorial problem is
    created: take for example an alkane chain: the carbon atoms match in one
    way. Then, for every :math:``CH_2`` group there are two, independent options
    for the hydrogen isomorphism. This would result in :math:``2^n`` subgraph
    isomorphisms for :math:``n`` carbon atoms.

    This means that the matches found will not be optimal for the hydrogens.
    This is acceptable, since hydrogrens are supposed to be equal. Let's say
    you have some sort of chiral atom with two hydrogens: it's not chiral and
    the hydrogen atoms are equal. Let's now say one of the two is a deuterium:
    in that case you should have a proper 'element' header, and the subgraph
    will be matched correctly.

    Parameters
    ----------
    reference : networkx.Graph
        The reference graph.
    residue : networkx.Graph
        The graph to match to ``reference``.
    Returns
    -------
    matches : list[dict]
        The matches found. The dictionaries have node indices of ``reference`` as
        keys and node indices of ``residue`` as values. Is an empty list if
        ``residue`` is not a subgraph of ``reference``.
    """
    # TODO: refactor this thing to accept node and edge compatibility checkers
    matches = []
#    H_idxs = [idx for idx in residue if residue.node[idx]['element'] == 'H']
    H_idxs = [idx for idx in residue if residue.degree(idx) == 1]
    heavy_res = nx.Graph(residue).copy()
    heavy_res.remove_nodes_from(H_idxs)

#    ref_H_idxs = [idx for idx in reference if reference.degree(idx) == 1]
#    heavy_ref = nx.Graph(reference).copy()
#    heavy_ref.remove_nodes_from(ref_H_idxs)
    # First, generate all the isomorphisms on heavy atoms. For each of these
    # we'll find *something* where the hydrogens match.
    GM = ElementGraphMatcher(reference, heavy_res)
    first_matches = list(GM.subgraph_isomorphisms_iter())
    for match in first_matches:
        reverse_match = {v: k for k, v in match.items()}
        for res_H_idx in H_idxs:
            # We know which parent atom this hydrogen is bound to, and we know
            # how it matches to the reference. We're going to find all the
            # neighboring atoms of the reference parent, and see if the name
            # of this hydrogen atom matches with any of those neighbors. If so,
            # we extend the match that way.
            # It should be noted that in exceptional cases where atomnames are
            # very wrong, this might cause a problem?
            res_neighbor = list(residue[res_H_idx].keys())[0]
            if res_neighbor not in reverse_match:
                continue
            ref_neighbor = reverse_match[res_neighbor]
            H_names = defaultdict(list)
            for idx in reference[ref_neighbor]:
                if reference.degree(idx) == 1:
                    H_names[reference.nodes[idx]['atomname']].append(idx)
            H_names = dict(H_names)
            res_H_name = residue.nodes[res_H_idx]['atomname']
            if res_H_name in H_names:
                if len(H_names[res_H_name]) != 1:
                    continue
                ref_H_idx = H_names[res_H_name][0]
                if ref_H_idx not in match and reference.nodes[ref_H_idx]['element'] == residue.nodes[res_H_idx]['element']:
                    reverse_match[res_H_idx] = ref_H_idx
                    match[ref_H_idx] = res_H_idx
        GM_large = ElementGraphMatcher(reference, residue)
        # Put the knowledge from the heavy atom isomorphism back in. Note that
        # ElementGraphMatched is modified to enable this and is no longer
        # re-entrant.
        # Indices in match do not have to be changed to account for interlaced
        # hydrogens: the node-indices in heavy_res and residue are the same.
        GM_large.core_1 = match  # pylint: disable=attribute-defined-outside-init
        GM_large.core_2 = reverse_match  # pylint: disable=attribute-defined-outside-init
        outcome = GM_large.subgraph_isomorphisms_iter()
        # Take just the first match found, otherwise it becomes a combinatorics
        # problem (consider an alkane chain). This is fine though, since
        # hydrogrens are supposed to be equal. Let's say you have some sort of
        # chiral atom with two hydrogens: It's not chiral. Let's now say one of
        # the two is a deuterium: in that case you should have a proper
        # 'element' header, and the subgraph will be matched correctly.
        # So worst case scenario we rename all hydrogens. This is acceptable
        # since they're equal.
        # And do islice since there may be none.
        # This will fail for e.g. oxygens which have degree 1 in the reference,
        # but are substituted with PTMs in the actual molecule. In that case
        # the atomnames might be flipped, and make PTM identification
        # troublesome. For example: C(=O)OH. That's why we extend the match to
        # include degree-1 nodes above.
        if match:
            matches.extend(itertools.islice(outcome, 1))
        else:
            matches.extend(outcome)
    matches = sorted(matches,
                     key=lambda m: rate_match(reference, residue, m),
                     reverse=True)
    return matches

class ElementGraphMatcher(nx.isomorphism.GraphMatcher):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        super().initialize()

    def initialize(self):
        return

    def semantic_feasibility(self, node1, node2):
        # TODO: implement (partial) wildcards
        elem1 = self.G1.node[node1]['element']
        elem2 = self.G2.node[node2]['element']
        return elem1 == elem2

def blockmodel(G, partitions, **attrs):
    """
    Analogous to networkx.blockmodel, but can deal with incomplete partitions,
    and assigns ``attrs`` to nodes.

    Parameters
    ----------
    G: networkx.Graph
        The graph to partition
    parititions: collections.abc.Iterable[collections.abc.Iterable]
        Each element contains the node indices that construct the new node.
    **attrs: dict[str, collections.abc.Iterable]
        Attributes to assign to new nodes. Attribute values are assigned to the
        new nodes in order.

    Returns
    -------
    networkx.Graph
        A new graph where every node is a subgraph as specified by partitions.
        Node attributes:

            :graph: Subgraph of constructing nodes.
            :nnodes: Number of nodes in ``graph``.
            :nedges: Number of edges in ``graph``.
            :density: Density of ``graph``.
            :attrs.keys(): As specified by ``**attrs``.
    """
    # TODO: Change this to use nx.quotient_graph.
    attrs = {key: list(val) for key, val in attrs.items()}
    CG_mol = nx.Graph()
    for bead_idx, idxs in enumerate(partitions):
        bd = G.subgraph(idxs)
        CG_mol.add_node(bead_idx)
        CG_mol.node[bead_idx]['graph'] = bd
        # TODO: CoM instead of CoG
#        CG_mol.node[bead_idx]['position'] = np.mean([bd.node[idx]['position'] for idx in bd], axis=0)
        for k, vals in attrs.items():
            CG_mol.node[bead_idx][k] = vals[bead_idx]

        CG_mol.node[bead_idx]['nnodes'] = bd.number_of_nodes()
        CG_mol.node[bead_idx]['nedges'] = bd.number_of_edges()
        CG_mol.node[bead_idx]['density'] = nx.density(bd)

    block_mapping = {}
    for n in CG_mol:
        nodes_in_block = CG_mol.node[n]['graph'].nodes()
        block_mapping.update(dict.fromkeys(nodes_in_block, n))

    for u, v, d in G.edges(data=True):
        try:
            bmu = block_mapping[u]
            bmv = block_mapping[v]
        except KeyError:
            # Atom not represented
            continue
        if bmu == bmv:  # no self loops
            continue
        # For graphs and digraphs add single weighted edge
        weight = d.get('weight', 1.0)  # default to 1 if no weight specified
        if CG_mol.has_edge(bmu, bmv):
            CG_mol[bmu][bmv]['weight'] += weight
        else:
            CG_mol.add_edge(bmu, bmv, weight=weight)
    return CG_mol

def rate_match(residue, bead, match):
    """
    A helper function which rates how well ``match`` describes the isomorphism
    between ``residue`` and ``bead`` based on the number of matching atomnames.

    Parameters
    ----------
    residue : networkx.Graph
        A graph. Required node attributes:

            :atomname: The name of an atom.

    bead : networkx.Graph
        A subgraph of ``residue`` where the isomorphism is described by ``match``.
        Required node attributes:

            :atomname: The name of an atom.

    Returns
    -------
    int
        The number of entries in match where the atomname in ``residue`` matches
        the atomname in ``bead``.
    """
    return sum(residue.node[rdx].get('atomname') == bead.node[bdx].get('atomname')
               for rdx, bdx in match.items())

def make_residue_graph(mol):
    """
    Creates a graph with one node per residue; as identified by the tuple
    (chain identifier, residue index, residue name).

    Parameters
    ----------
    mol: networkx.Graph
        The atomistic graph. Required node attributes:

            :chain: The chain identifier.
            :resid: The residue index.
            :resname: The residue name.

    Returns
    -------
    networkx.Graph
        A graph with one node per residue. Node attributes:

            :chain: The chain identifier.
            :graph: The atomistic subgraph.
            :density: The density of ``graph``.
            :nedges: The number of edges in ``graph``.
            :nnodes: The number of nodes in ``graph``.
            :resid: The residue index.
            :resname: The residue name.
            :atomname: The residue name.
    """
    def keyfunc(node_idx):
        return mol.node[node_idx]['chain'], mol.node[node_idx]['resid'], mol.node[node_idx]['resname']
    nodes = sorted(mol.node, key=keyfunc)
    keys = []
    grps = []
    for key, grp in itertools.groupby(nodes, keyfunc):
        keys.append(key)
        grps.append(list(grp))
    if keys:
        chain, resids, resnames = map(list, zip(*keys))
    else:
        chain, resids, resnames = [], [], []
    res_graph = blockmodel(mol, grps, chain=chain, resid=resids,
                           resname=resnames, atomname=resnames)
    return res_graph


# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
from collections import OrderedDict
import numpy as np
import astropy.units as u
from astropy.table import Table
from astropy.time import Time
from ..spectrum.utils import CountsPredictor
from ..stats.poisson import excess_error
from ..utils.scripts import make_path

__all__ = [
    'LightCurve',
    'LightCurveEstimator',
]

class LightCurve(object):
    """Lightcurve container.

    The lightcurve data is stored in ``table``.

    For now we only support times stored in MJD format!

    TODO: specification of format is work in progress
    See https://github.com/open-gamma-ray-astro/gamma-astro-data-formats/pull/61

    Usage: :ref:`time-lc`

    Parameters
    ----------
    table : `~astropy.table.Table`
        Table with lightcurve data
    """

    def __init__(self, table):
        self.table = table

    def __repr__(self):
        return '{}(len={})'.format(self.__class__.__name__, len(self.table))

    @property
    def time_scale(self):
        """Time scale (str).

        Taken from table "TIMESYS" header.
        Common values: "TT" or "UTC".
        Assumed default is "UTC".
        """
        return self.table.meta.get('TIMESYS', 'utc')

    @property
    def time_format(self):
        """Time format (str)."""
        return 'mjd'

    # @property
    # def time_ref(self):
    #     """Time reference (`~astropy.time.Time`)."""
    #     return time_ref_from_dict(self.table.meta)

    def _make_time(self, colname):
        val = self.table[colname].data
        scale = self.time_scale
        format = self.time_format
        return Time(val, scale=scale, format=format)

    @property
    def time(self):
        """Time (`~astropy.time.Time`)."""
        return self._make_time('time')

    @property
    def time_min(self):
        """Time bin start (`~astropy.time.Time`)."""
        return self._make_time('time_min')

    @property
    def time_max(self):
        """Time bin end (`~astropy.time.Time`)."""
        return self._make_time('time_max')

    @property
    def time_mid(self):
        """Time bin center (`~astropy.time.Time`).

        ::
            time_mid = time_min + 0.5 * time_delta
        """
        return self.time_min + 0.5 * self.time_delta

    @property
    def time_delta(self):
        """Time bin width (`~astropy.time.TimeDelta`).

        ::
            time_delta = time_max - time_min
        """
        return self.time_max - self.time_min

    @classmethod
    def read(cls, filename, **kwargs):
        """Read from file.

        Parameters
        ----------
        filename : str
            Filename
        kwargs : dict
            Keyword arguments passed to `astropy.table.Table.read`.
        """
        filename = make_path(filename)
        table = Table.read(str(filename), **kwargs)
        return cls(table=table)

    def write(self, filename, **kwargs):
        """Write to file.

        Parameters
        ----------
        filename : str
            Filename
        kwargs : dict
            Keyword arguments passed to `astropy.table.Table.write`.
        """
        filename = make_path(filename)
        self.table.write(str(filename), **kwargs)

    def compute_fvar(self):
        r"""Calculate the fractional excess variance.

        This method accesses the the ``FLUX`` and ``FLUX_ERR`` columns
        from the lightcurve data.

        The fractional excess variance :math:`F_{var}`, an intrinsic
        variability estimator, is given by

        .. math::
            F_{var} = \sqrt{\frac{S^{2} - \bar{\sigma^{2}}}{\bar{x}^{2}}}.

        It is the excess variance after accounting for the measurement errors
        on the light curve :math:`\sigma`. :math:`S` is the variance.

        Returns
        -------
        fvar, fvar_err : `~numpy.ndarray`
            Fractional excess variance.

        References
        ----------
        .. [Vaughan2003] "On characterizing the variability properties of X-ray light
           curves from active galaxies", Vaughan et al. (2003)
           http://adsabs.harvard.edu/abs/2003MNRAS.345.1271V
        """
        flux = self.table['flux'].data.astype('float64')
        flux_err = self.table['flux_err'].data.astype('float64')

        flux_mean = np.mean(flux)
        n_points = len(flux)

        s_square = np.sum((flux - flux_mean) ** 2) / (n_points - 1)
        sig_square = np.nansum(flux_err ** 2) / n_points
        fvar = np.sqrt(np.abs(s_square - sig_square)) / flux_mean

        sigxserr_a = np.sqrt(2 / n_points) * (sig_square / flux_mean) ** 2
        sigxserr_b = np.sqrt(sig_square / n_points) * (2 * fvar / flux_mean)
        sigxserr = np.sqrt(sigxserr_a ** 2 + sigxserr_b ** 2)
        fvar_err = sigxserr / (2 * fvar)

        return fvar, fvar_err

    def compute_chisq(self):
        """Calculate the chi-square test for `LightCurve`.

        Chisquare test is a variability estimator. It computes
        deviations from the expected value here mean value

        Returns
        -------
        ChiSq, P-value : tuple of float or `~numpy.ndarray`
            Tuple of Chi-square and P-value
        """
        import scipy.stats as stats
        flux = self.table['flux']
        yexp = np.mean(flux)
        yobs = flux.data
        chi2, pval = stats.chisquare(yobs, yexp)
        return chi2, pval

    def plot(self, ax=None):
        """Plot flux versus time.

        Parameters
        ----------
        ax : `~matplotlib.axes.Axes` or None, optional.
            The `~matplotlib.axes.Axes` object to be drawn on.
            If None, uses the current `~matplotlib.axes.Axes`.

        Returns
        -------
        ax : `~matplotlib.axes.Axes` or None, optional.
            The `~matplotlib.axes.Axes` object to be drawn on.
            If None, uses the current `~matplotlib.axes.Axes`.
        """
        import matplotlib.pyplot as plt
        ax = plt.gca() if ax is None else ax

        # TODO: Should we plot with normal time axis labels (ISO, not MJD)?

        x, xerr = self._get_plot_x()
        y, yerr = self._get_plot_y()

        ax.errorbar(x=x, y=y, xerr=xerr, yerr=yerr, linestyle="None")
        ax.scatter(x=x, y=y)
        ax.set_xlabel("Time (MJD)")
        ax.set_ylabel("Flux (cm-2 s-1)")

        return ax

    def _get_plot_x(self):
        try:
            x = self.time.mjd
        except KeyError:
            x = self.time_mid.mjd

        try:
            xerr = x - self.time_min.mjd, self.time_max.mjd - x
        except KeyError:
            xerr = None

        return x, xerr

    def _get_plot_y(self):
        y = self.table['flux'].quantity.to('cm-2 s-1').value

        if 'flux_errp' in self.table.colnames:
            yp = self.table['flux_errp'].quantity.to('cm-2 s-1').value
            yn = self.table['flux_errn'].quantity.to('cm-2 s-1').value
            yerr = yn, yp
        elif 'flux_err' in self.table.colnames:
            yerr = self.table['flux_err'].quantity.to('cm-2 s-1').value
        else:
            yerr = None

        return y, yerr

class LightCurveEstimator(object):
    """Light curve estimator.

    For a usage example see :gp-extra-notebook:`light_curve`.

    Parameters
    ----------
    spec_extract : `~gammapy.spectrum.SpectrumExtraction`
       Contains statistics, IRF and event lists
    """

    def __init__(self, spec_extract):
        self.obs_list = spec_extract.obs_list
        self.obs_spec = spec_extract.observations
        self.off_evt_list = self._get_off_evt_list(spec_extract)
        self.on_evt_list = self._get_on_evt_list(spec_extract)

    @staticmethod
    def _get_off_evt_list(spec_extract):
        """
        Returns list of OFF events for each observations
        """
        off_evt_list = []
        for bg in spec_extract.bkg_estimate:
            off_evt_list.append(bg.off_events)
        return off_evt_list

    @staticmethod
    def _get_on_evt_list(spec_extract):
        """
        Returns list of ON events for each observations
        """
        on_evt_list = []
        for obs in spec_extract.bkg_estimate:
            on_evt_list.append(obs.on_events)

        return on_evt_list

    @staticmethod
    def create_fixed_time_bin(time_step, spectrum_extraction):
        """Create time intervals of fixed size.

        Parameters
        ----------
        time_step : float
            Size of the light curve bins in seconds
        spectrum_extraction : `~gammapy.spectrum.SpectrumExtraction`
            Contains statistics, IRF and event lists

        Returns
        -------
        intervals : list of `~astropy.time.Time`
            List of time intervals
        """
        intervals = []
        time_start = Time(100000, format="mjd")
        time_end = Time(0, format="mjd")
        time_step = time_step / (24 * 3600)

        for obs in spectrum_extraction.obs_list:
            time_events = obs.events.time
            if time_start > time_events.min():
                time_start = time_events.min()
            if time_end < time_events.max():
                time_end = time_events.max()

        time = time_start.value
        while time < time_end.value:
            time += time_step
            intervals.append([
                Time(time - time_step, format="mjd"),
                Time(time, format="mjd"),
            ])
        return intervals

    def light_curve(self, time_intervals, spectral_model, energy_range):
        """Compute light curve.

        Implementation follows what is done in:
        http://adsabs.harvard.edu/abs/2010A%26A...520A..83H.

        To be discussed: assumption that threshold energy in the
        same in reco and true energy.

        Parameters
        ----------
        time_intervals : `list` of `~astropy.time.Time`
            List of time intervals
        spectral_model : `~gammapy.spectrum.models.SpectralModel`
            Spectral model
        energy_range : `~astropy.units.Quantity`
            True energy range to evaluate integrated flux (true energy)

        Returns
        -------
        lc : `~gammapy.time.LightCurve`
            Light curve
        """
        rows = []
        for time_interval in time_intervals:
            useinterval, row = self.compute_flux_point(time_interval, spectral_model, energy_range)
            if useinterval:
                rows.append(row)

        return self._make_lc_from_row_data(rows)

    @staticmethod
    def _make_lc_from_row_data(rows):
        table = Table()
        table['time_min'] = [_['time_min'].value for _ in rows]
        table['time_max'] = [_['time_max'].value for _ in rows]

        table['flux'] = [_['flux'].value for _ in rows] * u.Unit('1 / (s cm2)')
        table['flux_err'] = [_['flux_err'].value for _ in rows] * u.Unit('1 / (s cm2)')

        table['livetime'] = [_['livetime'].value for _ in rows] * u.s
        table['n_on'] = [_['n_on'] for _ in rows]
        table['n_off'] = [_['n_off'] for _ in rows]
        table['alpha'] = [_['alpha'] for _ in rows]
        table['measured_excess'] = [_['measured_excess'] for _ in rows]
        table['expected_excess'] = [_['expected_excess'].value for _ in rows]

        return LightCurve(table)

    def compute_flux_point(self, time_interval, spectral_model, energy_range):
        """Compute one flux point for one time interval.

        Parameters
        ----------
        time_interval : `~astropy.time.Time`
            Time interval (2-element array, or a tuple of Time objects)
        spectral_model : `~gammapy.spectrum.models.SpectralModel`
            Spectral model
        energy_range : `~astropy.units.Quantity`
            True energy range to evaluate integrated flux (true energy)

        Returns
        -------
        useinterval : bool
            Is True if the time_interval produce a valid flux point
        measurements : dict
            Dictionary with flux point measurement in the time interval
        """
        tmin, tmax = time_interval[0], time_interval[1]
        livetime = 0
        alpha_mean = 0.
        alpha_mean_backup = 0.
        measured_excess = 0
        predicted_excess = 0
        n_on = 0
        n_off = 0
        useinterval = False

        # Loop on observations
        for t_index, obs in enumerate(self.obs_list):

            spec = self.obs_spec[t_index]

            # discard observations not matching the time interval
            obs_start = obs.events.time[0]
            obs_stop = obs.events.time[-1]
            if (tmin < obs_start and tmax < obs_start) or (tmin > obs_stop):
                continue

            useinterval = True
            # get ON and OFF evt list
            off_evt = self.off_evt_list[t_index]
            on_evt = self.on_evt_list[t_index]

            # introduce the e_reco binning here, since it's also used
            # in the calculation of predicted counts
            e_reco = spec.e_reco
            emin = e_reco[e_reco.searchsorted(max(spec.lo_threshold, energy_range[0]))]
            emax = e_reco[e_reco.searchsorted(min(spec.hi_threshold, energy_range[1])) - 1]

            # compute ON events
            on = on_evt.select_energy([emin, emax])
            on = on.select_energy(energy_range)
            on = on.select_time([tmin, tmax])
            n_on_obs = len(on.table)

            # compute OFF events
            off = off_evt.select_energy([emin, emax])
            off = off.select_energy(energy_range)
            off = off.select_time([tmin, tmax])
            n_off_obs = len(off.table)

            # compute effective livetime (for the interval)
            if tmin >= obs_start and tmax <= obs_stop:
                # interval included in obs
                livetime_to_add = (tmax - tmin).to('s')
            elif tmin >= obs_start and tmax >= obs_stop:
                # interval min above tstart from obs
                livetime_to_add = (obs_stop - tmin).to('s')
            elif tmin <= obs_start and tmax <= obs_stop:
                # interval min below tstart from obs
                livetime_to_add = (tmax - obs_start).to('s')
            elif tmin <= obs_start and tmax >= obs_stop:
                # obs included in interval
                livetime_to_add = (obs_stop - obs_start).to('s')
            else:
                livetime_to_add = 0 * u.sec

            # Take into account dead time
            livetime_to_add *= (1. - obs.observation_dead_time_fraction)

            # Compute excess
            obs_measured_excess = n_on_obs - spec.alpha * n_off_obs

            # Compute the expected excess in the range given by the user
            # but must respect the energy threshold of the observation
            # (to match the energy range of the measured excess)
            # We use the effective livetime and the right energy threshold
            e_idx = np.where(np.logical_and.reduce(
                (e_reco >= spec.lo_threshold,  # threshold
                 e_reco <= spec.hi_threshold,  # threshold
                 e_reco >= energy_range[0],  # user
                 e_reco <= energy_range[-1])  # user
            ))[0]
            counts_predictor = CountsPredictor(
                livetime=livetime_to_add,
                aeff=spec.aeff,
                edisp=spec.edisp,
                model=spectral_model
            )
            counts_predictor.run()
            counts_predicted_excess = counts_predictor.npred.data.data[e_idx[:-1]]

            obs_predicted_excess = np.sum(counts_predicted_excess)

            # compute effective normalisation between ON/OFF (for the interval)
            livetime += livetime_to_add
            alpha_mean += spec.alpha * n_off_obs
            alpha_mean_backup += spec.alpha * livetime_to_add
            measured_excess += obs_measured_excess
            predicted_excess += obs_predicted_excess
            n_on += n_on_obs
            n_off += n_off_obs

        # Fill time interval information
        if useinterval:
            int_flux = spectral_model.integral(energy_range[0], energy_range[1])

            if n_off > 0.:
                alpha_mean /= n_off
            if livetime > 0.:
                alpha_mean_backup /= livetime
            if alpha_mean == 0.:  # use backup if necessary
                alpha_mean = alpha_mean_backup

            flux = measured_excess / predicted_excess.value
            flux *= int_flux
            flux_err = int_flux / predicted_excess.value
            # Gaussian errors, TODO: should be improved
            flux_err *= excess_error(n_on=n_on, n_off=n_off, alpha=alpha_mean)
        else:
            flux = 0
            flux_err = 0

        # Store measurements in a dict and return that
        return useinterval, OrderedDict([
            ('time_min', Time(tmin, format='mjd')),
            ('time_max', Time(tmax, format='mjd')),
            ('flux', flux * u.Unit('1 / (s cm2)')),
            ('flux_err', flux_err * u.Unit('1 / (s cm2)')),

            ('livetime', livetime * u.s),
            ('alpha', alpha_mean),
            ('n_on', n_on),
            ('n_off', n_off),
            ('measured_excess', measured_excess),
            ('expected_excess', predicted_excess),
        ])

import numpy,import scipy,import astropy,from astropy


"""
Retrieves either NZTA or NZS1170.5 code values
for the given locations
"""

from pathlib import Path
import argparse
import multiprocessing as mp
from typing import Sequence

import numpy as np
import pandas as pd

import sha_calc as sha
import gmhazard_calc as sc

DEFAULT_RETURN_PERIODS = np.array([20, 25, 50, 100, 250, 500, 1000, 2000, 2500])
DEFAULT_EXCEEDANCE_VALUES = 1 / DEFAULT_RETURN_PERIODS

DEFAULT_IMS = [
    "PGA",
    "pSA_0.01",
    "pSA_0.02",
    "pSA_0.03",
    "pSA_0.04",
    "pSA_0.05",
    "pSA_0.075",
    "pSA_0.1",
    "pSA_0.12",
    "pSA_0.15",
    "pSA_0.17",
    "pSA_0.2",
    "pSA_0.25",
    "pSA_0.3",
    "pSA_0.4",
    "pSA_0.5",
    "pSA_0.6",
    "pSA_0.7",
    "pSA_0.75",
    "pSA_0.8",
    "pSA_0.9",
    "pSA_1.0",
    "pSA_1.25",
    "pSA_1.5",
    "pSA_2.0",
    "pSA_2.5",
    "pSA_3.0",
    "pSA_4.0",
    "pSA_5.0",
    "pSA_6.0",
    "pSA_7.5",
    "pSA_10.0",
]

def main(
    input_data_ffp: str,
    output_dir: Path,
    nz_code_type: str,
    nzta_csv_ffp: Path = None,
    ims: Sequence[str] = DEFAULT_IMS,
    n_procs: int = 4,
):
    # Load the required data
    data_df = pd.read_csv(input_data_ffp)

    # Need an ensemble
    ens = sc.gm_data.Ensemble("v20p5emp")

    if nz_code_type == "NZS1170.5":
        with mp.Pool(n_procs) as pool:
            results = pool.starmap(
                _process_nzs1170p5_station,
                [
                    (
                        ens,
                        cur_row.lat,
                        cur_row.lon,
                        cur_row.vs30,
                        ims,
                        ix,
                        data_df.shape[0],
                    )
                    for ix, (cur_id, cur_row) in enumerate(data_df.iterrows())
                ],
            )

        # Extract and save
        grouped_sublists = list(zip(*results))
        np.save(
            str(output_dir / "NZS1170p5_im_values.npy"),
            np.stack(grouped_sublists[0], axis=0),
        )
        np.save(
            str(output_dir / "NZS1170p5_Z_values.npy"),
            np.stack(grouped_sublists[1], axis=0),
        )
        np.save(
            str(output_dir / "NZS1170p5_N_values.npy"),
            np.stack(grouped_sublists[2], axis=0),
        )
        np.save(
            str(output_dir / "NZS1170p5_R_values.npy"),
            np.stack(grouped_sublists[3], axis=0),
        )
        np.save(
            str(output_dir / "NZS1170p5_Ch_values.npy"),
            np.stack(grouped_sublists[4], axis=0),
        )

    elif nz_code_type == "NZTA":
        assert nzta_csv_ffp is not None, (
            "Path to the NZTA csv is required when " "computing NZTA code PGA values"
        )
        nzta_df = pd.read_csv(nzta_csv_ffp, header=0, index_col=0)

        with mp.Pool(n_procs) as pool:
            results = pool.starmap(
                _process_nzta_station,
                [
                    (
                        ens,
                        cur_row.lat,
                        cur_row.lon,
                        cur_row.vs30,
                        nzta_df,
                        ix,
                        data_df.shape[0],
                    )
                    for ix, (cur_id, cur_row) in enumerate(data_df.iterrows())
                ],
            )
            grouped_sublists = list(zip(*results))
            np.save(
                str(output_dir / "NZTA_PGA_values.npy"),
                np.stack(grouped_sublists[0], axis=0),
            )
            np.save(
                str(output_dir / "NZTA_town_index.npy"),
                np.stack(grouped_sublists[1], axis=0),
            )

def _process_nzta_station(
    ens: sc.gm_data.Ensemble,
    lat: float,
    lon: float,
    vs30: float,
    nzta_df: pd.DataFrame,
    ix: int,
    n_locs: int,
):
    print(f"Processing location {ix + 1}/{n_locs}")

    # Set result to nan if no vs30 values are available
    if np.isnan(vs30):
        return np.full(len(DEFAULT_EXCEEDANCE_VALUES), np.nan), np.nan

    site_info = sc.site.SiteInfo(f"site_{ix}", lat, lon, vs30)

    result = sc.nz_code.nzta_2018.run_ensemble_nzta(
        ens, site_info, exceedance_values=DEFAULT_EXCEEDANCE_VALUES
    )
    return (
        result.pga_values.loc[DEFAULT_EXCEEDANCE_VALUES].values,
        np.flatnonzero(nzta_df.index.values == result.nearest_town)[0],
    )

def _process_nzs1170p5_station(
    ens: sc.gm_data.Ensemble,
    lat: float,
    lon: float,
    vs30: float,
    ims: Sequence[str],
    ix: int,
    n_locs: int,
):
    print(f"Processing location {ix + 1}/{n_locs}")
    # Get the periods
    sa_periods = [0 if im == "PGA" else sc.utils.get_period_from_pSA(im) for im in ims]

    # Set result to nan if no vs30 values are available
    if np.isnan(vs30):
        return (
            np.full((len(sa_periods), len(DEFAULT_EXCEEDANCE_VALUES)), np.nan),
            np.nan,
            np.full(len(sa_periods), np.nan),
            np.full(len(DEFAULT_EXCEEDANCE_VALUES), np.nan),
            np.full(len(sa_periods), np.nan),
        )

    site_info = sc.site.SiteInfo(f"site_{ix}", lat, lon, vs30)

    distance = sc.nz_code.nzs1170p5.get_distance_from_site_info(ens, site_info)
    z_factor = float(sc.nz_code.nzs1170p5.ll2z((site_info.lon, site_info.lat)))
    soil_class = sc.nz_code.nzs1170p5.get_soil_class(site_info.vs30)

    results, R_values, N_values, Ch_values = [], [], [], []
    for cur_exceedance in DEFAULT_EXCEEDANCE_VALUES:
        cur_rp = 1 / cur_exceedance

        if cur_rp < 20 or cur_rp > 2500:
            raise NotImplementedError()
        else:
            C, Ch, R, N = sha.nzs1170p5_spectra(
                sa_periods, z_factor, cur_rp, distance, soil_class.value
            )
            results.append(C)
            R_values.append(R)
            N_values.append(N)
            Ch_values.append(Ch)

    return (
        np.stack(results, axis=1),
        z_factor,
        N_values[0],
        np.asarray(R_values),
        Ch_values[0],
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "input_data",
        type=str,
        help="Path to input csv, must contain columns lon, lat and vs30",
    )
    parser.add_argument("output_dir", type=Path, help="Output directory path")
    parser.add_argument(
        "nz_code_type",
        type=str,
        help="The NZCode for which to generate data",
        choices=["NZS1170.5", "NZTA"],
    )
    parser.add_argument(
        "--nzta_town_csv",
        type=Path,
        help="Path to the NZTA town csv, required when computing NZTA",
    )
    parser.add_argument(
        "--n_procs", type=int, help="Number of processes to use", default=4
    )

    args = parser.parse_args()

    main(
        args.input_data,
        args.output_dir,
        args.nz_code_type,
        nzta_csv_ffp=args.nzta_town_csv,
        n_procs=args.n_procs,
    )


import os, sys, glob, copy
import argparse
from pathlib import Path
from typing import Dict, Optional, Union, Type

import pandas as pd
import numpy as np
from functools import partial
from scipy.interpolate import RectBivariateSpline
from scipy.ndimage.filters import convolve

import astropy
from astropy.io import fits
from astropy import units as un
import lenstools
from lenstools import ConvergenceMap
import healpy as hp
#import pymaster as nmt

from astrild.simulation import Simulation
from astrild.rays.utils import Filters
from astrild.rays.skys.sky_array import SkyArray
from astrild.rays.skys.sky_namaster import SkyNamaster
from astrild.rays.skys.sky_utils import SkyUtils
from astrild.rays.skyio import SkyIO
from astrild.io import IO

dir_src = Path(__file__).parent.absolute()
default_config_file_ray = dir_src / "configs/ray_snapshot_info.h5"
c_light = 299792.458  # in km/s

class SkyHealpixWarning(BaseException):
    pass

class SkyHealpix:
    """
    The sky-map is constructed through multiple ray-tracing simulations
    run with RayRamses. This class analyzes the 2D map that contains
    the summes pertrurbations of each ray. It can prepare the data for the
    search of voids and peaks.

    Attributes:
        npix:
        theta:
        dirs:

    Methods:
        from_file:
        pdf:
        wl_peak_counts:
        add_field:
        add_galaxy_shape_noise:
        create_galaxy_shape_noise:
        add_cmb:
        create_cmb:
        add_mask:
        create_mask:
        convolution:
    """

    def __init__(
        self,
        skymap: np.ndarray,
        opening_angle: float,
        quantity: str,
        dirs: Dict[str, str],
        map_file: Optional[str] = None,
    ):
        self.data = {"orig": skymap}
        self._nside = hp.get_nside(skymap)
        self._npix = hp.nside2npix(self.nside)
        self.opening_angle = opening_angle
        self.quantity = quantity
        self.dirs = dirs
        self.map_file = map_file

    @classmethod
    def from_file(
        cls,
        map_file: str,
        opening_angle: float,
        quantity: str,
        dir_in: str,
        nside: Optional[int] = None,
        convert_unit: bool = True,
    ) -> "SkyHealpix":
        """
        Initialize class by reading the skymap data from pandas hdf5 file
        or numpy array.
        The file can be pointed at via map_filename or file_dsc.

        Args:
            map_filename:
                File path with which skymap pd.DataFrame can be loaded.
            file_dsc:
                Dictionary pointing to a file via {path, root, extension}.
                Use when multiple skymaps need to be loaded.
        """
        file_extension = map_file.split(".")[-1]
        assert file_extension in ["h5", "fits", "npy"], SkyHealpixWarning(
            f"The file formart {file_extension} is not supported."
        )
        if file_extension == "h5":
            map_df = pd.read_hdf(map_file, key="df")
            return cls.from_sky_dataframe(
                map_df,
                opening_angle,
                nside,
                quantity,
                dir_in,
                map_file,
                convert_unit,
            )
        elif file_extension == "fits":
            map_array = hp.read_map(map_file)
            return cls.from_sky_array(
                map_array, opening_angle, quantity, dir_in, map_file
            )
        elif file_extension == "npy":
            map_array = np.load(map_file)
            return cls.from_sky_array(
                map_array, opening_angle, quantity, dir_in, map_file
            )

    @classmethod
    def from_dataframe(
        cls,
        map_df: pd.DataFrame,
        opening_angle: float,
        nside: int,
        quantity: str,
        dir_in: str,
        map_file: str,
        convert_unit: bool = True,
    ) -> "SkyHealpix":
        """
        Initialize class by reading the skymap data from pandas DataFrame.

        Args:
            map_filename:
                File path with which skymap pd.DataFrame can be loaded.
            file_dsc:
                Dictionary pointing to a file via {path, root, extension}.
                Use when multiple skymaps need to be loaded.
        """
        if convert_unit:
            map_df = SkyUtils.convert_code_to_phy_units(quantity, map_df)
        map_array = SkyIO.transform_PandasDataFrame_to_Healpix(
            map_df, quantity, nside
        )
        return cls.from_sky_array(
            map_array, opening_angle, quantity, dir_in, map_file
        )

    @classmethod
    def from_array(
        cls,
        map_array: np.array,
        opening_angle: float,
        quantity: str,
        dir_in: str,
        map_file: Optional[str] = None,
    ) -> "SkyHealpix":
        """
        Initialize class by reading the skymap data from np.ndarray.

        Args:
            map_filename:
                File path with which skymap pd.DataFrame can be loaded.
            file_dsc:
                Dictionary pointing to a file via {path, root, extension}.
                Use when multiple skymaps need to be loaded.
        """
        dirs = {"sim": dir_in}
        map_array = hp.ma(map_array)  # mask out bad values (e.g Nan)
        return cls(map_array, opening_angle, quantity, dirs, map_file)

    @classmethod
    def from_Cl_file(
        cls,
        cl_file: str,
        quantity: str,
        dir_in: str,
        nside: int,
        lmax: int = 3000,
        opening_angle: int = 41253,
        key: Optional[str] = None,
        rnd_seed: Optional[int] = None,
    ) -> "SkyHealpix":
        """
        Args:
            cl_file:
            quantity:
            dir_in:
            nside:
            lmax:
            opening_angle:
            key:
        """
        np.random.seed(rnd_seed)
        file_extension = cl_file.split(".")[-1]
        assert file_extension in ["npz", "npy"], SkyHealpixWarning(
            f"The file formart {file_extension} is not supported."
        )
        if file_extension == "npy":
            cl_array = np.load(cl_file)
        elif file_extension == "npz":
            cl_array = np.load(cl_file)[key]
        return cls.from_Cl_array(
            cl_array, quantity, nside, lmax, opening_angle, dir_in, cl_file
        )

    @classmethod
    def from_Cl_array(
        cls,
        cl_array: np.array,
        quantity: str,
        nside: int,
        lmax: int,
        opening_angle: int = 41253,
        dir_in: Optional[str] = None,
        cl_file: Optional[str] = None,
        rnd_seed: Optional[int] = None,
    ) -> "SkyHealpix":
        """
        Args:
            opening_angle:
                [deg^2]
        """
        np.random.seed(rnd_seed)
        dirs = {"sim": dir_in}
        map_array = hp.sphtfunc.synfast(cl_array, nside=nside, lmax=lmax)
        return cls(map_array, opening_angle, quantity, dirs, cl_file)

    @property
    def nside(self):
        return self._nside

    @property
    def npix(self):
        return self._npix

    def to_skyarray(
        self,
        npix: int,
        opening_angle: float,
        of: str = "orig",
    ) -> Type[SkyArray]:
        """
        Args:
            lonra: , [deg]
            latra: , [deg]
        """
        lonra = latra = [0, opening_angle]
        cart_proj = hp.projector.CartesianProj(
            lonra=lonra,
            latra=latra,
            xsize=npix,
            ysize=npix,
        )
        map_array = cart_proj.projmap(
            self.data[of],
            rot=(0., 0.),
            vec2pix_func=partial(hp.vec2pix, self._nside),
        )
        return SkyArray.from_array(
            map_array,
            opening_angle=opening_angle,
            quantity=self.quantity,
            dir_in=self.dirs["sim"],
        )

    def create_cmb(
        self,
        filepath_cl: str,
        lmax: float = 3e3,
        nside: Optional[int] = None,
        rnd_seed: Optional[int] = None,
    ) -> None:
        """
        Cosmig Microwave Background (CMB) on partial-sky map,
        for which the flat-sky approximation holds (ell > 10).

        Args:
            filepath_cl:
                angular power spectrum of CMB
            nside:
                Nr. of pixels per edge of the output full-sky map
            rnd_seed:
                Fix random seed, for reproducability.

        Returns:
            cmb_map:
        """
        if nside is None:
            nside = self._nside
        cl_cmb = np.load(filepath_cl)
        np.random.seed(rnd_seed)
        self.data["cmb"] = hp.sphtfunc.synfast(cmb_tt, nside=nside, lmax=lmax)

    def sum_of_maps(self, map1: str, map2: str) -> None:
        self.data[f"{map1}_{map2}"] = self.data[map1] + self.data[map2]

    def arithmetic_operation_with(
        self, skymap: np.array, on: str, operation: np
    ) -> None:
        """
        Use Numpy function to perform one of the basic arithmetic operations:
            add, substract, multiply, divide
        of two Healpix fields.

        Args:
            skymap:
            on:
            operation:
        """
        _pixel_idx = np.arange(self._npix)
        _unmasked_pixel_idx = _pixel_idx[~self.data[on].mask]
        self.data[on].data[_unmasked_pixel_idx] = operation(
            self.data[on].data[_unmasked_pixel_idx], skymap[_unmasked_pixel_idx]
        )

    def add_mask(self, on: str, theta: Optional[float] = None) -> None:
        """
        """
        if ("mask" not in self.data.keys()) or (theta != self.mask_theta):
            self.create_mask(theta)
        self.data[on + "_mask"] = hp.ma(copy.deepcopy(self.data[on]))
        self.data[on + "_mask"].mask = self.data["mask"]

    def create_mask(self, theta: float) -> None:
        """
        Mask out unobserved patches of the full-sky.

        Args:
            theta:
                Edge length of the square field-of-view [deg]
            nside:
                Nr. of pixels per edge of the output full-sky map
        """
        print("create_mask", theta)
        # angular positions of all healpix pixels
        _pixel_index = np.arange(self._npix)
        pixel_theta, pixel_phi = hp.pix2ang(self._nside, _pixel_index)
        # range of ra and dec of the field-of-view
        ras = np.array([90 - theta / 2, 90 + theta / 2]) * np.pi / 180  # [rad]
        decs = np.array([theta / 2, 360 - theta / 2]) * np.pi / 180  # [rad]
        # create mask
        mask = np.zeros(self._npix, dtype=np.bool)
        mask[pixel_theta < ras[0]] = 1
        mask[pixel_theta > ras[1]] = 1
        mask[(decs[0] < pixel_phi) & (pixel_phi < decs[1])] = 1
        self.data["mask"] = mask
        self.mask_theta = theta

    def rotate(
        self, theta: float, phi: float, which: str, rtn: bool = False
    ) -> np.ndarray:
        """
        Rotate the full-sky.

        Args:
            theta and phi:
                Rotational angles [deg]
            which:
                Identify which map in self.data should be rotated.
        """
        if (theta is None) or (phi is None):
            theta = np.random.random() * 180
            phi = np.random.random() * 180

        # Get theta, phi for non-rotated map
        _pixel_index = np.arange(self._npix)
        t, p = hp.pix2ang(self._nside, _pixel_index)
        # Define a rotator
        r = hp.Rotator(deg=True, rot=[theta, phi])
        # Get theta, phi under rotated co-ordinates
        trot, prot = r(t, p)
        # Interpolate map onto these co-ordinates
        _skymap = copy.deepcopy(self.data[which])
        _skymap = hp.get_interp_val(_skymap, trot, prot)
        if rtn is True:
            return _skymap
        else:
            self.data[which] = _skymap

    #def to_skynamaster(self, apodization_width: float = 2.5):
    #    """
    #    Transform SkyHealpix to SkyNamaster.
    #    """
    #    _skyarray = copy.deepcopy(self.data["orig"])
    #    _skyarray[~self.data["mask"]] = hp.UNSEEN
    #    _mask = copy.deepcopy(self.data["mask"]) * 1
    #    #_mask = nmt.mask_apodization(_mask, apodization_width, apotype="C2")
    #    #print("------------->", _mask.shape)
    #    #self.data["nm"] = nmt.NmtField(_mask, [_skyarray])
    #    # return SkyNamaster.from_array(
    #    #    sky_field,
    #    #    self._npix,
    #    #    self.opening_angle,
    #    #    self.quantity,
    #    #    self.dirs,
    #    #    self.map_file,
    #    # )

import numpy,from scipy,import astropy,from astropy