# Licensed under a 3-clause BSD style license - see LICENSE.rst
import logging
import numpy as np
from itertools import combinations
from astropy.coordinates import Angle, SkyOffsetFrame
from astropy.table import Table
import astropy.units as u
from gammapy.irf import EDispMap, PSFMap, FoVAlignment
from gammapy.maps import Map, RegionGeom, RegionNDMap
from gammapy.modeling.models import PowerLawSpectralModel
from gammapy.stats import WStatCountsStatistic
from gammapy.utils.coordinates import sky_to_fov

__all__ = [
    "make_counts_rad_max",
    "make_edisp_kernel_map",
    "make_edisp_map",
    "make_map_background_irf",
    "make_map_exposure_true_energy",
    "make_psf_map",
    "make_theta_squared_table",
]

log = logging.getLogger(__name__)

def make_map_exposure_true_energy(
    pointing, livetime, aeff, geom, use_region_center=True
):
    """Compute exposure map.

    This map has a true energy axis, the exposure is not combined
    with energy dispersion.

    Parameters
    ----------
    pointing : `~astropy.coordinates.SkyCoord`
        Pointing direction
    livetime : `~astropy.units.Quantity`
        Livetime
    aeff : `~gammapy.irf.EffectiveAreaTable2D`
        Effective area
    geom : `~gammapy.maps.WcsGeom`
        Map geometry (must have an energy axis)
    use_region_center: bool
        If geom is a RegionGeom, whether to just
        consider the values at the region center
        or the instead the average over the whole region

    Returns
    -------
    map : `~gammapy.maps.WcsNDMap`
        Exposure map
    """
    if not use_region_center:
        coords, weights = geom.get_wcs_coord_and_weights()
    else:
        coords, weights = geom.get_coord(sparse=True), None

    offset = coords.skycoord.separation(pointing)
    exposure = aeff.evaluate(offset=offset, energy_true=coords["energy_true"])

    data = (exposure * livetime).to("m2 s")
    meta = {"livetime": livetime, "is_pointlike": aeff.is_pointlike}

    if not use_region_center:
        data = np.average(data, axis=1, weights=weights)

    return Map.from_geom(geom=geom, data=data.value, unit=data.unit, meta=meta)

def _map_spectrum_weight(map, spectrum=None):
    """Weight a map with a spectrum.

    This requires map to have an "energy" axis.
    The weights are normalised so that they sum to 1.
    The mean and unit of the output image is the same as of the input cube.

    At the moment this is used to get a weighted exposure image.

    Parameters
    ----------
    map : `~gammapy.maps.Map`
        Input map with an "energy" axis.
    spectrum : `~gammapy.modeling.models.SpectralModel`
        Spectral model to compute the weights.
        Default is power-law with spectral index of 2.

    Returns
    -------
    map_weighted : `~gammapy.maps.Map`
        Weighted image
    """
    if spectrum is None:
        spectrum = PowerLawSpectralModel(index=2.0)

    # Compute weights vector
    energy_edges = map.geom.axes["energy_true"].edges
    weights = spectrum.integral(
        energy_min=energy_edges[:-1], energy_max=energy_edges[1:]
    )
    weights /= weights.sum()
    shape = np.ones(len(map.geom.data_shape))
    shape[0] = -1
    return map * weights.reshape(shape.astype(int))

def make_map_background_irf(
    pointing, ontime, bkg, geom, oversampling=None, use_region_center=True
):
    """Compute background map from background IRFs.

    Parameters
    ----------
    pointing : `~gammapy.data.FixedPointingInfo` or `~astropy.coordinates.SkyCoord`
        Observation pointing

        - If a `~gammapy.data.FixedPointingInfo` is passed, FOV coordinates are properly computed.
        - If a `~astropy.coordinates.SkyCoord` is passed, FOV frame rotation is not taken into account.
    ontime : `~astropy.units.Quantity`
        Observation ontime. i.e. not corrected for deadtime
        see https://gamma-astro-data-formats.readthedocs.io/en/stable/irfs/full_enclosure/bkg/index.html#notes)
    bkg : `~gammapy.irf.Background3D`
        Background rate model
    geom : `~gammapy.maps.WcsGeom`
        Reference geometry
    oversampling: int
        Oversampling factor in energy, used for the background model evaluation.
    use_region_center: bool
        If geom is a RegionGeom, whether to just
        consider the values at the region center
        or the instead the sum over the whole region

    Returns
    -------
    background : `~gammapy.maps.WcsNDMap`
        Background predicted counts sky cube in reco energy
    """
    # TODO:
    #  This implementation can be improved in two ways:
    #  1. Create equal time intervals between TSTART and TSTOP and sum up the
    #  background IRF for each interval. This is instead of multiplying by
    #  the total ontime. This then handles the rotation of the FoV.
    #  2. Use the pointing table (does not currently exist in CTA files) to
    #  obtain the RA DEC and time for each interval. This then considers that
    #  the pointing might change slightly over the observation duration

    # Get altaz coords for map
    if oversampling is not None:
        geom = geom.upsample(factor=oversampling, axis_name="energy")

    coords = {"energy": geom.axes["energy"].edges.reshape((-1, 1, 1))}

    if not use_region_center:
        image_geom = geom.to_wcs_geom().to_image()
        region_coord, weights = geom.get_wcs_coord_and_weights()
        idx = image_geom.coord_to_idx(region_coord)
        sky_coord = region_coord.skycoord
        d_omega = image_geom.solid_angle().T[idx]
    else:
        image_geom = geom.to_image()
        map_coord = image_geom.get_coord()
        sky_coord = map_coord.skycoord
        d_omega = image_geom.solid_angle()

    if bkg.has_offset_axis:
        coords["offset"] = sky_coord.separation(pointing)
    else:
        if bkg.fov_alignment == FoVAlignment.ALTAZ:
            altaz_coord = sky_coord.transform_to(pointing.altaz_frame)

            # Compute FOV coordinates of map relative to pointing
            fov_lon, fov_lat = sky_to_fov(
                altaz_coord.az, altaz_coord.alt, pointing.altaz.az, pointing.altaz.alt
            )
        elif bkg.fov_alignment == FoVAlignment.RADEC:
            # Create OffsetFrame
            frame = SkyOffsetFrame(origin=pointing.radec)
            pseudo_fov_coord = sky_coord.transform_to(frame)
            fov_lon = pseudo_fov_coord.lon
            fov_lat = pseudo_fov_coord.lat
        else:
            raise ValueError(
                f"Unsupported background coordinate system: {bkg.fov_alignment!r}"
            )

        coords["fov_lon"] = fov_lon
        coords["fov_lat"] = fov_lat

    bkg_de = bkg.integrate_log_log(**coords, axis_name="energy")
    data = (bkg_de * d_omega * ontime).to_value("")

    if not use_region_center:
        data = np.sum(weights * data, axis=2)

    bkg_map = Map.from_geom(geom, data=data)

    if oversampling is not None:
        bkg_map = bkg_map.downsample(factor=oversampling, axis_name="energy")

    return bkg_map

def make_psf_map(psf, pointing, geom, exposure_map=None):
    """Make a psf map for a single observation

    Expected axes : rad and true energy in this specific order
    The name of the rad MapAxis is expected to be 'rad'

    Parameters
    ----------
    psf : `~gammapy.irf.PSF3D`
        the PSF IRF
    pointing : `~astropy.coordinates.SkyCoord`
        the pointing direction
    geom : `~gammapy.maps.Geom`
        the map geom to be used. It provides the target geometry.
        rad and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        the associated exposure map.
        default is None

    Returns
    -------
    psfmap : `~gammapy.irf.PSFMap`
        the resulting PSF map
    """
    coords = geom.get_coord(sparse=True)

    # Compute separations with pointing position
    offset = coords.skycoord.separation(pointing)

    # Compute PSF values
    data = psf.evaluate(
        energy_true=coords["energy_true"],
        offset=offset,
        rad=coords["rad"],
    )

    # Create Map and fill relevant entries
    psf_map = Map.from_geom(geom, data=data.value, unit=data.unit)
    psf_map.normalize(axis_name="rad")
    return PSFMap(psf_map, exposure_map)

def make_edisp_map(edisp, pointing, geom, exposure_map=None, use_region_center=True):
    """Make a edisp map for a single observation

    Expected axes : migra and true energy in this specific order
    The name of the migra MapAxis is expected to be 'migra'

    Parameters
    ----------
    edisp : `~gammapy.irf.EnergyDispersion2D`
        the 2D Energy Dispersion IRF
    pointing : `~astropy.coordinates.SkyCoord`
        the pointing direction
    geom : `~gammapy.maps.Geom`
        the map geom to be used. It provides the target geometry.
        migra and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        the associated exposure map.
        default is None
    use_region_center: Bool
        If geom is a RegionGeom, whether to just
        consider the values at the region center
        or the instead the average over the whole region

    Returns
    -------
    edispmap : `~gammapy.irf.EDispMap`
        the resulting EDisp map
    """
    # Compute separations with pointing position
    if not use_region_center:
        coords, weights = geom.get_wcs_coord_and_weights()
    else:
        coords, weights = geom.get_coord(sparse=True), None

    offset = coords.skycoord.separation(pointing)

    # Compute EDisp values
    data = edisp.evaluate(
        offset=offset,
        energy_true=coords["energy_true"],
        migra=coords["migra"],
    )

    if not use_region_center:
        data = np.average(data, axis=2, weights=weights)

    # Create Map and fill relevant entries
    edisp_map = Map.from_geom(geom, data=data.to_value(""), unit="")
    edisp_map.normalize(axis_name="migra")
    return EDispMap(edisp_map, exposure_map)

def make_edisp_kernel_map(
    edisp, pointing, geom, exposure_map=None, use_region_center=True
):
    """Make a edisp kernel map for a single observation

    Expected axes : (reco) energy and true energy in this specific order
    The name of the reco energy MapAxis is expected to be 'energy'.
    The name of the true energy MapAxis is expected to be 'energy_true'.

    Parameters
    ----------
    edisp : `~gammapy.irf.EnergyDispersion2D`
        the 2D Energy Dispersion IRF
    pointing : `~astropy.coordinates.SkyCoord`
        the pointing direction
    geom : `~gammapy.maps.Geom`
        the map geom to be used. It provides the target geometry.
        energy and true energy axes should be given in this specific order.
    exposure_map : `~gammapy.maps.Map`, optional
        the associated exposure map.
        default is None
    use_region_center: Bool
        If geom is a RegionGeom, whether to just
        consider the values at the region center
        or the instead the average over the whole region

    Returns
    -------
    edispmap : `~gammapy.irf.EDispKernelMap`
        the resulting EDispKernel map
    """
    # Use EnergyDispersion2D migra axis.
    migra_axis = edisp.axes["migra"]

    # Create temporary EDispMap Geom
    new_geom = geom.to_image().to_cube([migra_axis, geom.axes["energy_true"]])

    edisp_map = make_edisp_map(
        edisp, pointing, new_geom, exposure_map, use_region_center
    )

    return edisp_map.to_edisp_kernel_map(geom.axes["energy"])

def make_theta_squared_table(
    observations, theta_squared_axis, position, position_off=None
):
    """Make theta squared distribution in the same FoV for a list of `Observation`
    objects.

    The ON theta2 profile is computed from a given distribution, on_position.
    By default, the OFF theta2 profile is extracted from a mirror position
    radially symmetric in the FOV to pos_on.

    The ON and OFF regions are assumed to be of the same size, so the normalisation
    factor between both region alpha = 1.

    Parameters
    ----------
    observations: `~gammapy.data.Observations`
        List of observations
    theta_squared_axis : `~gammapy.maps.geom.MapAxis`
        Axis of edges of the theta2 bin used to compute the distribution
    position : `~astropy.coordinates.SkyCoord`
        Position from which the on theta^2 distribution is computed
    position_off : `astropy.coordinates.SkyCoord`
        Position from which the OFF theta^2 distribution is computed.
        Default: reflected position w.r.t. to the pointing position

    Returns
    -------
    table : `~astropy.table.Table`
        Table containing the on counts, the off counts, acceptance, off acceptance and alpha
        for each theta squared bin.
    """
    if not theta_squared_axis.edges.unit.is_equivalent("deg2"):
        raise ValueError("The theta2 axis should be equivalent to deg2")

    table = Table()

    table["theta2_min"] = theta_squared_axis.edges[:-1]
    table["theta2_max"] = theta_squared_axis.edges[1:]
    table["counts"] = 0
    table["counts_off"] = 0
    table["acceptance"] = 0.0
    table["acceptance_off"] = 0.0

    alpha_tot = np.zeros(len(table))
    livetime_tot = 0

    create_off = position_off is None
    for observation in observations:
        separation = position.separation(observation.events.radec)
        counts, _ = np.histogram(separation ** 2, theta_squared_axis.edges)
        table["counts"] += counts

        if create_off:
            # Estimate the position of the mirror position
            pos_angle = observation.pointing_radec.position_angle(position)
            sep_angle = observation.pointing_radec.separation(position)
            position_off = observation.pointing_radec.directional_offset_by(
                pos_angle + Angle(np.pi, "rad"), sep_angle
            )

        # Angular distance of the events from the mirror position
        separation_off = position_off.separation(observation.events.radec)

        # Extract the ON and OFF theta2 distribution from the two positions.
        counts_off, _ = np.histogram(separation_off ** 2, theta_squared_axis.edges)
        table["counts_off"] += counts_off

        # Normalisation between ON and OFF is one
        acceptance = np.ones(theta_squared_axis.nbin)
        acceptance_off = np.ones(theta_squared_axis.nbin)

        table["acceptance"] += acceptance
        table["acceptance_off"] += acceptance_off
        alpha = acceptance / acceptance_off
        alpha_tot += alpha * observation.observation_live_time_duration.to_value("s")
        livetime_tot += observation.observation_live_time_duration.to_value("s")

    alpha_tot /= livetime_tot
    table["alpha"] = alpha_tot

    stat = WStatCountsStatistic(table["counts"], table["counts_off"], table["alpha"])
    table["excess"] = stat.n_sig
    table["sqrt_ts"] = stat.sqrt_ts
    table["excess_errn"] = stat.compute_errn()
    table["excess_errp"] = stat.compute_errp()

    table.meta["ON_RA"] = position.icrs.ra
    table.meta["ON_DEC"] = position.icrs.dec
    return table

def make_counts_rad_max(geom, rad_max, events):
    """Extract the counts using for the ON region size the values in the
    `RAD_MAX_2D` table.

    Parameters
    ----------
    geom : `~gammapy.maps.RegionGeom`
        reference map geom
    rad_max : `~gammapy.irf.RadMax2D`
        the RAD_MAX_2D table IRF
    events : `~gammapy.data.EventList`
        event list to be used to compute the ON counts

    Returns
    -------
    counts : `~gammapy.maps.RegionNDMap`
        Counts vs estimated energy extracted from the ON region.
    """
    selected_events = events.select_rad_max(
        rad_max=rad_max, position=geom.region.center
    )

    counts = Map.from_geom(geom=geom)
    counts.fill_events(selected_events)
    return counts

def are_regions_overlapping_rad_max(regions, rad_max, offset, e_min, e_max):
    """
    Calculate pair-wise separations between all regions and compare with rad_max
    to find overlaps.
    """
    separations = u.Quantity([
        a.center.separation(b.center)
        for a, b in combinations(regions, 2)
    ])

    rad_max_at_offset = rad_max.evaluate(offset=offset)
    # do not check bins outside of energy range
    edges_min = rad_max.axes['energy'].edges_min
    edges_max = rad_max.axes['energy'].edges_max
    # to be sure all possible values are included, we check
    # for the *upper* energy bin to be larger than e_min and the *lower* edge
    # to be larger than e_max
    mask = (edges_max >= e_min) & (edges_min <= e_max)
    rad_max_at_offset = rad_max_at_offset[mask]

    return np.any(separations[np.newaxis, :] < (2 * rad_max_at_offset))

def make_counts_off_rad_max(
    on_geom,
    rad_max,
    events,
    region_finder,
    exclusion_mask=None,
):
    """Extract the OFF counts and the ON / OFF acceptance considering for the
    sizes of the ON and OFF regions the values in the `RAD_MAX_2D` table.
    Per each estimated energy bin a `ReflectedRegionsFinder` is defined to
    search for the OFF regions.

    Parameters
    ----------
    geom: `~gammapy.maps.RegionGeom`
        reference map geom for the on region
    rad_max: `~gammapy.irf.RadMax2D`
        the RAD_MAX_2D table IRF
    events: `~gammapy.data.EventList`
        event list to be used to compute the OFF counts
    region_finder: `~gammapy.makers.background.reflected.RegionFinder`

    Returns
    -------
    counts_off : `~gammapy.maps.RegionNDMap`
        OFF Counts vs estimated energy extracted from the ON region.
    acceptance_off : `~gammapy.maps.RegionNDMap`
        ratio of the acceptances of the OFF to ON regions.
    """

    off_regions, wcs = region_finder.run(
        center=events.pointing_radec,
        region=on_geom.region,
        exclusion_mask=exclusion_mask,
    )

    if len(off_regions) == 0:
        log.warning("RegionsFinder returned no regions")
        # counts_off=None, acceptance_off=0
        return None, RegionNDMap.from_geom(on_geom, data=0)

    # check for overlap
    energy_axis = on_geom.axes["energy"]
    offset = on_geom.region.center.separation(events.pointing_radec)
    e_min, e_max = energy_axis.edges[[0, -1]]
    regions = [on_geom.region] + off_regions
    if are_regions_overlapping_rad_max(regions, rad_max, offset, e_min, e_max):
        log.warning("Found overlapping on/off regions, choose less off regions")
        # counts_off=None, acceptance_off=0
        return None, RegionNDMap.from_geom(on_geom, data=0)

    off_region_geom = RegionGeom.from_regions(
        regions=off_regions,
        axes=[energy_axis],
        wcs=wcs,
    )

    counts_off = RegionNDMap.from_geom(geom=off_region_geom)
    acceptance_off = RegionNDMap.from_geom(
        geom=off_region_geom,
        data=np.full(energy_axis.nbin, len(off_regions))
    )

    for off_region in off_regions:
        selected_events = events.select_rad_max(
            rad_max=rad_max, position=off_region.center
        )
        counts_off.fill_events(selected_events)

    return counts_off, acceptance_off

import numpy,import astropy,from astropy


#!/usr/bin/env python

from __future__ import division
import numpy as np
import scipy.spatial.distance as ssd
import h5py, sys

import pycuda.driver as drv
import pycuda.autoinit
from pycuda import gpuarray
from scikits.cuda import linalg
linalg.init()

from tps import tps_kernel_matrix, tps_eval
from transformations import unit_boxify
from culinalg_exts import dot_batch_nocheck, get_gpu_ptrs, m_dot_batch
from precompute import downsample_cloud, batch_get_sol_params
from cuda_funcs import init_prob_nm, norm_prob_nm, get_targ_pts, check_cuda_err, \
    fill_mat, reset_cuda, sq_diffs
from registration import registration_cost as cpu_registration_cost
from defaults import N_ITER_CHEAP, EM_ITER_CHEAP, DEFAULT_LAMBDA, MAX_CLD_SIZE, \
    DATA_DIM, DS_SIZE, N_STREAMS, DEFAULT_NORM_ITERS, BEND_COEF_DIGITS

import IPython as ipy
from pdb import pm, set_trace
import time

class Globals:
    sync = False
    streams = []
    for i in range(N_STREAMS):
        streams.append(drv.Stream())

def get_stream(i):
    return Globals.streams[i % N_STREAMS]

def sync(override = False):
    if Globals.sync or override:
        check_cuda_err()

def loglinspace(a,b,n):
    "n numbers between a to b (inclusive) with constant ratio between consecutive numbers"
    return np.exp(np.linspace(np.log(a),np.log(b),n))

def gpu_pad(x, shape, dtype=np.float32):
    (m, n) = x.shape
    if m > shape[0] or n > shape[1]:
        raise ValueError("Cannot Pad Beyond Normal Dimension")
    x_new = np.zeros(shape, dtype=dtype)
    x_new[:m, :n] = x
    return gpuarray.to_gpu(x_new)

class GPUContext(object):
    """
    Class to contain GPU arrays
    """
    def __init__(self, bend_coefs = None):
        if bend_coefs is None:
            lambda_init, lambda_final = DEFAULT_LAMBDA
            bend_coefs = np.around(loglinspace(lambda_init, lambda_final, N_ITER_CHEAP),
                                    BEND_COEF_DIGITS)
        self.bend_coefs = bend_coefs
        self.ptrs_valid = False
        self.N = 0

        self.tps_params     = []
        self.tps_param_ptrs = None
        self.trans_d        = []
        self.trans_d_ptrs   = None
        self.lin_dd         = []
        self.lin_dd_ptrs    = None
        self.w_nd           = []
        self.w_nd_ptrs      = None
        self.ident_mats     = [gpuarray.to_gpu(np.eye(MAX_CLD_SIZE, dtype=np.float32))]
        self.ident_ptrs      = None

        """
        TPS PARAM FORMAT
        [      np.zeros(DATA_DIM)      ]   [  trans_d  ]   [1 x d]
        [       np.eye(DATA_DIM)       ] = [  lin_dd   ] = [d x d]
        [np.zeros((np.zeros, DATA_DIM))]   [   w_nd    ]   [n x d]
        """
        self.default_tps_params = gpuarray.zeros((DATA_DIM + 1 + MAX_CLD_SIZE, DATA_DIM), np.float32)
        self.default_tps_params[1:DATA_DIM+1, :].set(np.eye(DATA_DIM, dtype=np.float32))

        self.proj_mats       = dict([(b, []) for b in bend_coefs])
        self.proj_mat_ptrs   = dict([(b, None) for b in bend_coefs])
        self.offset_mats     = dict([(b, []) for b in bend_coefs])
        self.offset_mat_ptrs = dict([(b, None) for b in bend_coefs])

        self.pts             = []
        self.pt_ptrs         = None
        self.pt_t_pt         = []
        self.pt_t_pt_ptrs    = None
        self.kernels         = []
        self.kernel_ptrs     = None
        self.pts_w           = []
        self.pt_w_ptrs       = None
        self.pts_t           = []
        self.pts_t_ptrs      = None
        self.dims            = []
        self.dims_gpu        = None

        self.warp_err         = None
        self.bend_res         = []
        self_bend_res_ptrs    = None

        self.corr_cm          = []
        self.corr_cm_ptrs     = None
        self.corr_rm          = []
        self.corr_rm_ptrs     = None
        self.r_coefs         = []
        self.r_coef_ptrs     = None
        self.c_coefs_rn      = []
        self.c_coef_rn_ptrs  = None
        self.c_coefs_cn      = []
        self.c_coef_cn_ptrs  = None
        self.seg_names        = []

    def reset_tps_params(self):
        """
        sets the tps params to be identity
        """
        for p in self.tps_params:
            drv.memcpy_dtod_async(p.gpudata, self.default_tps_params.gpudata, p.nbytes)
    def set_tps_params(self, vals):
        for d, s in zip(self.tps_params, vals):
            drv.memcpy_dtod_async(d.gpudata, s.gpudata, d.nbytes)

    def reset_warp_err(self):
        self.warp_err.fill(0)

    def check_cld(self, cloud_xyz):
        if cloud_xyz.dtype != np.float32:
            raise TypeError("only single precision operations supported")
        if cloud_xyz.shape[0] > MAX_CLD_SIZE:
            raise ValueError("cloud size exceeds {}".format(MAX_CLD_SIZE))
        if cloud_xyz.shape[1] != DATA_DIM:
            raise ValueError("point cloud must have cumn dimension {}".format(DATA_DIM))
    # @profile
    def get_sol_params(self, cld):
        self.check_cld(cld)
        K = tps_kernel_matrix(cld)
        proj_mats   = {}
        offset_mats = {}
        (proj_mats_arr, _), (offset_mats_arr, _) = batch_get_sol_params(cld, K, self.bend_coefs)
        for i, b in enumerate(self.bend_coefs):
            proj_mats[b]   = proj_mats_arr[i]
            offset_mats[b] = offset_mats_arr[i]
        return proj_mats, offset_mats, K

    def add_cld(self, name, proj_mats, offset_mats, cloud_xyz, kernel, update_ptrs = False):
        """
        adds a new cloud to our context for batch processing
        """
        self.check_cld(cloud_xyz)
        self.ptrs_valid = False
        self.N += 1
        self.seg_names.append(name)
        self.tps_params.append(self.default_tps_params.copy())
        self.trans_d.append(self.tps_params[-1][0, :])
        self.lin_dd.append(self.tps_params[-1][1:DATA_DIM+1, :])
        self.w_nd.append(self.tps_params[-1][DATA_DIM + 1:, :])
        n = cloud_xyz.shape[0]

        for b in self.bend_coefs:
            proj_mat   = proj_mats[b]
            offset_mat = offset_mats[b]
            self.proj_mats[b].append(gpu_pad(proj_mat, (MAX_CLD_SIZE + DATA_DIM + 1, MAX_CLD_SIZE)))

            if offset_mat.shape != (n + DATA_DIM + 1, DATA_DIM):
                raise ValueError("Offset Matrix has incorrect dimension")
            self.offset_mats[b].append(gpu_pad(offset_mat, (MAX_CLD_SIZE + DATA_DIM + 1, DATA_DIM)))

        if n > MAX_CLD_SIZE or cloud_xyz.shape[1] != DATA_DIM:
            raise ValueError("cloud_xyz has incorrect dimension")
        self.pts.append(gpu_pad(cloud_xyz, (MAX_CLD_SIZE, DATA_DIM)))
        self.pt_t_pt.append(gpu_pad(cloud_xyz.T.dot(cloud_xyz), (MAX_CLD_SIZE, MAX_CLD_SIZE)))
        if kernel.shape != (n, n):
            raise ValueError("dimension mismatch b/t kernel and cloud")
        self.kernels.append(gpu_pad(kernel, (MAX_CLD_SIZE, MAX_CLD_SIZE)))
        self.dims.append(n)

        self.pts_w.append(gpuarray.zeros_like(self.pts[-1]))
        self.pts_t.append(gpuarray.zeros_like(self.pts[-1]))
        self.corr_cm.append(gpuarray.zeros((MAX_CLD_SIZE, MAX_CLD_SIZE), np.float32))
        self.corr_rm.append(gpuarray.zeros((MAX_CLD_SIZE, MAX_CLD_SIZE), np.float32))
        self.r_coefs.append(gpuarray.zeros((MAX_CLD_SIZE, 1), np.float32))
        self.c_coefs_rn.append(gpuarray.zeros((MAX_CLD_SIZE, 1), np.float32))
        self.c_coefs_cn.append(gpuarray.zeros((MAX_CLD_SIZE, 1), np.float32))

        if update_ptrs:
            self.update_ptrs()

    def update_ptrs(self):
        self.tps_param_ptrs = get_gpu_ptrs(self.tps_params)
        self.trans_d_ptrs   = get_gpu_ptrs(self.trans_d)
        self.lin_dd_ptrs    = get_gpu_ptrs(self.lin_dd)
        self.w_nd_ptrs      = get_gpu_ptrs(self.w_nd)

        for b in self.bend_coefs:
            self.proj_mat_ptrs[b]   = get_gpu_ptrs(self.proj_mats[b])
            self.offset_mat_ptrs[b] = get_gpu_ptrs(self.offset_mats[b])

        self.pt_ptrs         = get_gpu_ptrs(self.pts)
        self.pt_t_pt_ptrs    = get_gpu_ptrs(self.pt_t_pt)
        self.kernel_ptrs     = get_gpu_ptrs(self.kernels)
        self.ident_ptrs      = get_gpu_ptrs([self.ident_mats[0] for _ in range(self.N)])
        self.pt_w_ptrs       = get_gpu_ptrs(self.pts_w)
        self.pt_t_ptrs       = get_gpu_ptrs(self.pts_t)
        self.corr_cm_ptrs    = get_gpu_ptrs(self.corr_cm)
        self.corr_rm_ptrs    = get_gpu_ptrs(self.corr_rm)
        self.r_coef_ptrs    = get_gpu_ptrs(self.r_coefs)
        self.c_coef_rn_ptrs = get_gpu_ptrs(self.c_coefs_rn)
        self.c_coef_cn_ptrs = get_gpu_ptrs(self.c_coefs_cn)
        ## temporary space used for bend/warp cost computations
        self.warp_err        = gpuarray.zeros((self.N, MAX_CLD_SIZE), np.float32)
        self.bend_res_mat    = gpuarray.zeros((DATA_DIM * self.N, DATA_DIM), np.float32)
        self.bend_res        =[self.bend_res_mat[i*DATA_DIM:(i+1)*DATA_DIM] for i in range(self.N)]
        self.bend_res_ptrs   = get_gpu_ptrs(self.bend_res)

        self.dims_gpu = gpuarray.to_gpu(np.array(self.dims, dtype=np.int32))
        self.ptrs_valid = True

    def read_h5(self, fname):
        f = h5py.File(fname, 'r')
        for seg_name, seg_info in f.iteritems():
            if 'inv' not in seg_info:
                raise KeyError("H5 File does not have precomputed values")
            seg_info = seg_info['inv']

            proj_mats   = {}
            offset_mats = {}
            for b in self.bend_coefs:
                k = str(b)
                if k not in seg_info:
                    raise KeyError("H5 File {} bend coeficient {}".format(seg_name, k))
                proj_mats[b] = seg_info[k]['proj_mat'][:]
                offset_mats[b] = seg_info[k]['offset_mat'][:]

            ds_key    = 'DS_SIZE_{}'.format(DS_SIZE)
            cloud_xyz = seg_info[ds_key]['scaled_cloud_xyz'][:]
            kernel    = seg_info[ds_key]['scaled_K_nn'][:]
            self.add_cld(seg_name, proj_mats, offset_mats, cloud_xyz, kernel)
        f.close()
        self.update_ptrs()

    # @profile
    def setup_tgt_ctx(self, cloud_xyz):
        """
        returns a GPUContext where all the clouds are cloud_xyz
        and matched in length with this contex

        assumes cloud_xyz is already downsampled and scaled
        """
        tgt_ctx = TgtContext(self)
        tgt_ctx.set_cld(cloud_xyz)
        return tgt_ctx

    # @profile
    def transform_points(self):
        """
        computes the warp of self.pts under the current tps params
        """
        fill_mat(self.pt_w_ptrs, self.trans_d_ptrs, self.dims_gpu, self.N)
        dot_batch_nocheck(self.pts,         self.lin_dd,      self.pts_w,
                          self.pt_ptrs,     self.lin_dd_ptrs, self.pt_w_ptrs)
        dot_batch_nocheck(self.kernels,     self.w_nd,        self.pts_w,
                          self.kernel_ptrs, self.w_nd_ptrs,   self.pt_w_ptrs)
        sync()

    # @profile
    def get_target_points(self, other, outlierprior=1e-1, outlierfrac=1e-2, outliercutoff=1e-2,
                          T = 5e-3, norm_iters = DEFAULT_NORM_ITERS):
        """
        computes the target points for self and other
        using the current warped points for both
        """
        init_prob_nm(self.pt_ptrs, other.pt_ptrs,
                     self.pt_w_ptrs, other.pt_w_ptrs,
                     self.dims_gpu, other.dims_gpu,
                     self.N, outlierprior, outlierfrac, T,
                     self.corr_cm_ptrs, self.corr_rm_ptrs)
        sync()
        norm_prob_nm(self.corr_cm_ptrs, self.corr_rm_ptrs,
                     self.dims_gpu, other.dims_gpu, self.N, outlierfrac, norm_iters,
                     self.r_coef_ptrs, self.c_coef_rn_ptrs, self.c_coef_cn_ptrs)
        sync()
        get_targ_pts(self.pt_ptrs, other.pt_ptrs,
                     self.pt_w_ptrs, other.pt_w_ptrs,
                     self.corr_cm_ptrs, self.corr_rm_ptrs,
                     self.r_coef_ptrs, self.c_coef_rn_ptrs, self.c_coef_cn_ptrs,
                     self.dims_gpu, other.dims_gpu,
                     outliercutoff, self.N,
                     self.pt_t_ptrs, other.pt_t_ptrs)
        sync()

    # @profile
    def update_transform(self, b):
        """
        computes the TPS associated with the current target pts
        """
        self.set_tps_params(self.offset_mats[b])
        dot_batch_nocheck(self.proj_mats[b],     self.pts_t,     self.tps_params,
                          self.proj_mat_ptrs[b], self.pt_t_ptrs, self.tps_param_ptrs)
        sync()
    # @profile
    def mapping_cost(self, other, bend_coef=DEFAULT_LAMBDA[1], outlierprior=1e-1, outlierfrac=1e-2,
                       outliercutoff=1e-2,  T = 5e-3, norm_iters = DEFAULT_NORM_ITERS):
        """
        computes the error in the current mapping
        assumes that the target points have already been filled
        """
        self.transform_points()
        other.transform_points()
        sums = []
        sq_diffs(self.pt_w_ptrs, self.pt_t_ptrs, self.warp_err, self.N, True)
        sq_diffs(other.pt_w_ptrs, other.pt_t_ptrs, self.warp_err, self.N, False)
        warp_err = self.warp_err.get()
        return np.sum(warp_err, axis=1)
    # @profile
    def bending_cost(self, b=DEFAULT_LAMBDA[1]):
        ## b * w_nd' * K * w_nd
        ## use pts_w as temporary storage
        dot_batch_nocheck(self.kernels,     self.w_nd,      self.pts_w,
                          self.kernel_ptrs, self.w_nd_ptrs, self.pt_w_ptrs,
                          b = 0)

        dot_batch_nocheck(self.pts_w,     self.w_nd,      self.bend_res,
                          self.pt_w_ptrs, self.w_nd_ptrs, self.bend_res_ptrs,
                          transa='T', b = 0)
        bend_res = self.bend_res_mat.get()
        return b * np.array([np.trace(bend_res[i*DATA_DIM:(i+1)*DATA_DIM]) for i in range(self.N)])
    # @profile
    def bidir_tps_cost(self, other, bend_coef=1, outlierprior=1e-1, outlierfrac=1e-2,
                       outliercutoff=1e-2,  T = 5e-3, norm_iters = DEFAULT_NORM_ITERS):
        self.reset_warp_err()
        mapping_err  = self.mapping_cost(other, outlierprior, outlierfrac, outliercutoff, T, norm_iters)
        bending_cost = self.bending_cost(bend_coef)
        bending_cost += other.bending_cost(bend_coef)
        return mapping_err + bending_cost

    """
    testing for custom kernels
    """

    def test_mapping_cost(self, other, bend_coef=DEFAULT_LAMBDA[1], outlierprior=1e-1, outlierfrac=1e-2,
                       outliercutoff=1e-2,  T = 5e-3, norm_iters = DEFAULT_NORM_ITERS):
        mapping_err = self.mapping_cost(other, outlierprior, outlierfrac, outliercutoff, T, norm_iters)
        for i in range(self.N):
            ## compute error for 0 on cpu
            s_gpu = mapping_err[i]
            s_cpu = np.float32(0)
            xt = self.pts_t[i].get()
            xw = self.pts_w[i].get()

            yt = other.pts_t[i].get()
            yw = other.pts_w[i].get()

            ##use the trace b/c then numpy will use float32s all the way
            s_cpu += np.trace(xt.T.dot(xt) + xw.T.dot(xw) - 2 * xw.T.dot(xt))
            s_cpu += np.trace(yt.T.dot(yt) + yw.T.dot(yw) - 2 * yw.T.dot(yt))

            if not np.isclose(s_cpu, s_gpu, atol=1e-4):
                ## high err tolerance is b/c of difference in cpu and gpu precision?
                print "cpu and gpu sum sq differences differ!!!"
                ipy.embed()
                sys.exit(1)

    def test_bending_cost(self, other, bend_coef=DEFAULT_LAMBDA[1], outlierprior=1e-1, outlierfrac=1e-2,
                       outliercutoff=1e-2,  T = 5e-3, norm_iters = DEFAULT_NORM_ITERS):
        self.get_target_points(other, outlierprior, outlierfrac, outliercutoff,  T, norm_iters)
        self.update_transform(bend_coef)
        bending_costs = self.bending_cost(bend_coef)
        for i in range(self.N):
            c_gpu = bending_costs[i]
            k_nn = self.kernels[i].get()
            w_nd = self.w_nd[i].get()
            c_cpu = np.float32(0)
            for d in range(DATA_DIM):
                r = np.dot(k_nn, w_nd[:, d]).astype(np.float32)
                r = np.float32(np.dot(w_nd[:, d], r))
                c_cpu += r
            c_cpu *= np.float32(bend_coef)
            if np.abs(c_cpu - c_gpu) > 1e-4:
                ## high err tolerance is b/c of difference in cpu and gpu precision?
                print "cpu and gpu bend costs differ!!!"
                ipy.embed()
                sys.exit(1)

    def test_init_corr(self, other, T = 5e-3, outlierprior=1e-1, outlierfrac=1e-2, outliercutoff=1e-2, ):
        import scipy.spatial.distance as ssd
        import sys
        self.transform_points()
        other.transform_points()
        init_prob_nm(self.pt_ptrs, other.pt_ptrs,
                     self.pt_w_ptrs, other.pt_w_ptrs,
                     self.dims_gpu, other.dims_gpu,
                     self.N, outlierprior, outlierfrac, T,
                     self.corr_cm_ptrs, self.corr_rm_ptrs)
        gpu_corr_rm = self.corr_rm[0].get()
        gpu_corr_rm = gpu_corr_rm.flatten()[:(self.dims[0] + 1) * (other.dims[0] + 1)].reshape(self.dims[0]+1, other.dims[0]+1)
        s_pt_w = self.pts_w[0].get()
        s_pt   = self.pts[0].get()
        o_pt_w = other.pts_w[0].get()
        o_pt   = other.pts[0].get()

        d1 = ssd.cdist(s_pt_w, o_pt, 'euclidean')
        d2 = ssd.cdist(s_pt, o_pt_w, 'euclidean')

        p_nm = np.exp( -(d1 + d2) / (2 * T))

        for i in range(self.dims[0]):
            for j in range(other.dims[0]):
                if abs(p_nm[i, j] - gpu_corr_rm[i, j]) > 1e-7:
                    print "INIT CORR MATRICES DIFFERENT"
                    print i, j, p_nm[i, j], gpu_corr_rm[i, j]
                    ipy.embed()
                    sys.exit(1)

    def test_norm_corr(self, other, T = 5e-3, outlierprior=1e-1, outlierfrac=1e-2, outliercutoff=1e-2, norm_iters = DEFAULT_NORM_ITERS):
        import sys
        self.transform_points()
        other.transform_points()
        init_prob_nm(self.pt_ptrs, other.pt_ptrs,
                     self.pt_w_ptrs, other.pt_w_ptrs,
                     self.dims_gpu, other.dims_gpu,
                     self.N, outlierprior, outlierfrac, T,
                     self.corr_cm_ptrs, self.corr_rm_ptrs)
        n, m  = self.dims[0], other.dims[0]
        init_corr = self.corr_rm[0].get()
        init_corr = init_corr.flatten()[:(n + 1) * (m + 1)].reshape(n+1, m+1).astype(np.float32)

        a_N = np.ones((n+1),dtype = np.float32)
        a_N[n] = m*outlierfrac
        b_M = np.ones((m+1), dtype = np.float32)
        b_M[m] = n*outlierfrac

        old_r_coefs = np.ones(n+1, dtype=np.float32)
        old_c_coefs = np.ones(m+1, dtype=np.float32)
        for n_iter in range(1, norm_iters):
            init_prob_nm(self.pt_ptrs, other.pt_ptrs,
                         self.pt_w_ptrs, other.pt_w_ptrs,
                         self.dims_gpu, other.dims_gpu,
                         self.N, outlierprior, outlierfrac, T,
                         self.corr_cm_ptrs, self.corr_rm_ptrs)

            norm_prob_nm(self.corr_cm_ptrs, self.corr_rm_ptrs,
                         self.dims_gpu, other.dims_gpu, self.N, outlierfrac, n_iter,
                         self.r_coef_ptrs, self.c_coef_rn_ptrs, self.c_coef_cn_ptrs)
            new_r_coefs = a_N/init_corr.dot(old_c_coefs[:m+1])
            new_c_coefs = b_M/new_r_coefs[:n+1].dot(init_corr)
            gpu_c_coefs = self.c_coefs_cn[0].get().flatten()[:m + 1].reshape(m+1)
            gpu_r_coefs = self.r_coefs[0].get().flatten()[:n + 1].reshape(n+1)
            if not np.allclose(new_r_coefs, gpu_r_coefs):
                print "row coeficients don't match", n_iter
                ipy.embed()
                sys.exit(1)
            if not np.allclose(new_c_coefs, gpu_c_coefs):
                print "column coeficients don't match", n_iter
                ipy.embed()
                sys.exit(1)
            # old_r_coefs = gpu_r_coefs
            # old_c_coefs = gpu_c_coefs
            old_r_coefs = new_r_coefs
            old_c_coefs = new_c_coefs

    def test_get_targ(self, other, T = 5e-3, outlierprior=1e-1, outlierfrac=1e-2, outliercutoff=1e-2, norm_iters = DEFAULT_NORM_ITERS):
        self.transform_points()
        other.transform_points()
        init_prob_nm(self.pt_ptrs, other.pt_ptrs,
                     self.pt_w_ptrs, other.pt_w_ptrs,
                     self.dims_gpu, other.dims_gpu,
                     self.N, outlierprior, outlierfrac, T,
                     self.corr_cm_ptrs, self.corr_rm_ptrs)
        norm_prob_nm(self.corr_cm_ptrs, self.corr_rm_ptrs,
                     self.dims_gpu, other.dims_gpu, self.N, outlierfrac, norm_iters,
                     self.r_coef_ptrs, self.c_coef_rn_ptrs, self.c_coef_cn_ptrs)
        get_targ_pts(self.pt_ptrs, other.pt_ptrs,
                     self.pt_w_ptrs, other.pt_w_ptrs,
                     self.corr_cm_ptrs, self.corr_rm_ptrs,
                     self.r_coef_ptrs, self.c_coef_rn_ptrs, self.c_coef_cn_ptrs,
                     self.dims_gpu, other.dims_gpu,
                     outliercutoff, self.N,
                     self.pt_t_ptrs, other.pt_t_ptrs)
        n, m = self.dims[0], other.dims[0]
        x = self.pts[0].get()[:n]
        xw = self.pts_w[0].get()[:n]
        xt = self.pts_t[0].get()[:n]
        y = other.pts[0].get()[:m]
        yw = other.pts_w[0].get()[:m]
        yt = other.pts_t[0].get()[:m]

        init_corr = self.corr_rm[0].get()
        init_corr = init_corr.flatten()[:(n + 1) * (m + 1)].reshape(n+1, m+1).astype(np.float32)
        gpu_c_cn_coefs = self.c_coefs_cn[0].get().flatten()[:m + 1].reshape(m+1)
        gpu_c_rn_coefs = self.c_coefs_rn[0].get().flatten()[:m + 1].reshape(m+1)
        gpu_r_coefs = self.r_coefs[0].get().flatten()[:n + 1].reshape(n+1)

        rn_corr = (init_corr * gpu_c_rn_coefs[None, :]) * gpu_r_coefs[:, None]
        cn_corr = (init_corr * gpu_c_cn_coefs[None, :]) * gpu_r_coefs[:, None]
        rn_corr = rn_corr[:n, :m]
        cn_corr = cn_corr[:n, :m]

        wt_n = rn_corr.sum(axis=1)
        inlier = wt_n > outliercutoff
        xtarg = np.empty((n, DATA_DIM))
        xtarg[inlier, :] = rn_corr.dot(y)[inlier, :]
        xtarg[~inlier, :] = xw[~inlier, :]

        if not np.allclose(xtarg, xt):
            print "xt values differ"
            ipy.embed()
            sys.exit(1)

        wt_m = cn_corr.sum(axis=0)
        inliner = wt_m > outliercutoff
        ytarg = np.empty((m, DATA_DIM))
        ytarg[inlier, :] = cn_corr.T.dot(x)[inliner, :]
        ytarg[~inlier, :] = yw[~inlier, :]
        if not np.allclose(ytarg, yt):
            print "yt values differ"
            ipy.embed()
            sys.exit(1)

    def unit_test(self, other):
        print "running basic unit tests"
        self.test_init_corr(other)
        self.test_norm_corr(other)
        self.test_get_targ(other)
        self.test_mapping_cost(other)
        self.test_bending_cost(other)
        print "UNIT TESTS PASSED"

class TgtContext(GPUContext):
    """
    specialized class to handle the case where we are
    mapping to a single target cloud --> only allocate GPU Memory once
    """
    def __init__(self, src_ctx):
        GPUContext.__init__(self, src_ctx.bend_coefs)
        self.src_ctx = src_ctx
        ## just setup with 0's
        tgt_cld = np.zeros((MAX_CLD_SIZE, DATA_DIM), np.float32)
        proj_mats = dict([(b, np.zeros((MAX_CLD_SIZE + DATA_DIM + 1, MAX_CLD_SIZE), np.float32))
                          for b in self.bend_coefs])
        offset_mats = dict([(b, np.zeros((MAX_CLD_SIZE + DATA_DIM + 1, DATA_DIM), np.float32))
                            for b in self.bend_coefs])
        tgt_K = np.zeros((MAX_CLD_SIZE, MAX_CLD_SIZE), np.float32)
        for n in src_ctx.seg_names:
            name = "{}_tgt".format(n)
            GPUContext.add_cld(self, name, proj_mats, offset_mats, tgt_cld, tgt_K)
        GPUContext.update_ptrs(self)
    def add_cld(self, name, proj_mats, offset_mats, cloud_xyz, kernel, update_ptrs = False):
        raise NotImplementedError("not implemented for TgtConext")
    def update_ptrs(self):
        raise NotImplementedError("not implemented for TgtConext")
    # @profile
    def set_cld(self, cld):
        """
        sets the cloud for this appropriately
        won't allocate any new memory
        """
        proj_mats, offset_mats, K = self.get_sol_params(cld)
        K_gpu = gpu_pad(K, (MAX_CLD_SIZE, MAX_CLD_SIZE))
        cld_gpu = gpu_pad(cld, (MAX_CLD_SIZE, DATA_DIM))
        self.pts         = [cld_gpu for _ in range(self.N)]
        self.kernels     = [K_gpu for _ in range(self.N)]
        proj_mats_gpu    = dict([(b, gpu_pad(p.get(), (MAX_CLD_SIZE + DATA_DIM + 1, MAX_CLD_SIZE)))
                                 for b, p in proj_mats.iteritems()])
        self.proj_mats   = dict([(b, [p for _ in range(self.N)])
                                 for b, p in proj_mats_gpu.iteritems()])
        offset_mats_gpu  = dict([(b, gpu_pad(p.get(), (MAX_CLD_SIZE + DATA_DIM + 1, DATA_DIM)))
                                 for b, p in offset_mats.iteritems()])
        self.offset_mats = dict([(b, [p for _ in range(self.N)])
                                 for b, p in offset_mats_gpu.iteritems()])
        self.dims        = [cld.shape[0]]

        self.pt_ptrs.fill(int(self.pts[0].gpudata))
        self.kernel_ptrs.fill(int(self.kernels[0].gpudata))
        self.dims_gpu.fill(self.dims[0])
        for b in self.bend_coefs:
            self.proj_mat_ptrs[b].fill(int(self.proj_mats[b][0].gpudata))
            self.offset_mat_ptrs[b].fill(int(self.offset_mats[b][0].gpudata))

def check_transform_pts(ctx, i = 0):
    import scikits.cuda.linalg as la
    n = ctx.dims[i]
    w_nd = ctx.w_nd[i].get()[:n]
    lin_dd = ctx.lin_dd[i].get()
    trans_d = ctx.trans_d[i].get()
    k_nn = ctx.kernels[i].get()[:n, :n].reshape(n, n).copy()
    x_nd = ctx.pts[i].get()[:n]
    xw_nd = ctx.pts_w[i].get()[:n]

    _k_gpu = gpuarray.to_gpu(k_nn)
    _x_gpu = gpuarray.to_gpu(x_nd)
    _lin_gpu = gpuarray.to_gpu(lin_dd)
    _trans_gpu = gpuarray.to_gpu(trans_d)
    _w_gpu = gpuarray.to_gpu(w_nd)

    fill_mat(ctx.pt_w_ptrs, ctx.trans_d_ptrs, ctx.dims_gpu, ctx.N)
    dot_batch_nocheck(ctx.pts,         ctx.lin_dd,      ctx.pts_w,
                      ctx.pt_ptrs,     ctx.lin_dd_ptrs, ctx.pt_w_ptrs)

    xw_nd = ctx.pts_w[i].get()[:n]
    cpu_xw_nd = np.dot(x_nd, lin_dd) + trans_d[None, :]
    # assert np.allclose(xw_nd, cpu_xw_nd)

    dot_batch_nocheck(ctx.kernels,     ctx.w_nd,        ctx.pts_w,
                      ctx.kernel_ptrs, ctx.w_nd_ptrs,   ctx.pt_w_ptrs)
    xw_nd = ctx.pts_w[i].get()[:n]
    cpu_xw_nd = cpu_xw_nd + np.dot(k_nn, w_nd)
    # print "w_nd\n", w_nd[:3], np.max(w_nd)
    # print "lin_dd\n", lin_dd[:3]
    # print "trans_d\n", trans_d
    # print "k_nn\n", k_nn[:3, :3]
    # print "x_nd\n", x_nd[:3, :3]
    # print cpu_xw_nd[:3]
    if not(np.allclose(xw_nd, cpu_xw_nd) ):
        print "k dot w_nd is difference on cpu and gpu"
        k_dot_w = np.dot(k_nn, w_nd)
        k_gpu = [gpuarray.to_gpu(k_nn)]
        w_gpu = [gpuarray.to_gpu(w_nd)]
        res_gpu = [gpuarray.zeros((n, DATA_DIM), np.float32)]
        k_ptrs = get_gpu_ptrs(k_gpu)
        w_ptrs = get_gpu_ptrs(w_gpu)
        res_ptrs = get_gpu_ptrs(res_gpu)
        dot_batch_nocheck(k_gpu, w_gpu, res_gpu, k_ptrs, w_ptrs, res_ptrs)
        res = res_gpu[0].get()
        single_gpu = la.dot(_k_gpu, _w_gpu)
        print "retry success {}".format(np.allclose(res, k_dot_w))
        print "gpu success {}".format(np.allclose(single_gpu.get(), res))
        assert np.allclose(single_gpu.get(), res)
        raw_input("go?")

def check_update(ctx, b):
    ctx.tps_params[0] = ctx.default_tps_params.copy()
    ctx.update_ptrs()
    xt = ctx.pts_t[0].get()
    p_mat = ctx.proj_mats[b][0].get()
    o_mat = ctx.offset_mats[b][0].get()
    true_res = np.dot(p_mat, xt) + o_mat
    ctx.set_tps_params(ctx.offset_mats[b])
    o_gpu = ctx.tps_params[0].get()
    if not np.allclose(o_gpu, o_mat):
        print "setting tps params failed"
        diff = np.abs(o_mat - o_gpu)
        nz = np.nonzero(diff)
        print nz
        ipy.embed()
        sys.exit(1)
    ctx.update_transform(b)
    p1 = ctx.tps_params[0].get()
    if not np.allclose(true_res, p1):
        print "p1 and true res differ"
        print p1[:3]
        diff = np.abs(p1 - true_res)
        print np.max(diff)
        amax = np.argmax(diff)
        print amax
        nz = np.nonzero(diff)
        print nz[0]
        ipy.embed()
        sys.exit(1)

# @profile
def batch_tps_rpm_bij(src_ctx, tgt_ctx, T_init = 1e-1, T_final = 5e-3,
                      outlierfrac = 1e-2, outlierprior = 1e-1, outliercutoff = 1e-2, em_iter = EM_ITER_CHEAP):
    """
    computes tps rpm for the clouds in src and tgt in batch
    TODO: Fill out comment cleanly
    """
    ##TODO: add check to ensure that src_ctx and tgt_ctx are formatted properly
    n_iter = len(src_ctx.bend_coefs)
    T_vals = loglinspace(T_init, T_final, n_iter)

    src_ctx.reset_tps_params()
    tgt_ctx.reset_tps_params()
    for i, b in enumerate(src_ctx.bend_coefs):
        T = T_vals[i]
        for _ in range(em_iter):
            src_ctx.transform_points()
            tgt_ctx.transform_points()
            src_ctx.get_target_points(tgt_ctx, outlierprior, outlierfrac, outliercutoff, T)
            src_ctx.update_transform(b)
            # check_update(src_ctx, b)
            tgt_ctx.update_transform(b)
    return src_ctx.bidir_tps_cost(tgt_ctx)

def test_batch_tps_rpm_bij(src_ctx, tgt_ctx, T_init = 1e-1, T_final = 5e-3,
                           outlierfrac = 1e-2, outlierprior = 1e-1, outliercutoff = .5, em_iter = EM_ITER_CHEAP,
                           test_ind = 0):
    from transformations import ThinPlateSpline, set_ThinPlateSpline
    import tps
    n_iter = len(src_ctx.bend_coefs)
    T_vals = loglinspace(T_init, T_final, n_iter)

    x_nd = src_ctx.pts[test_ind].get()[:src_ctx.dims[test_ind]]
    y_md = tgt_ctx.pts[0].get()[:tgt_ctx.dims[0]]
    (n, d) = x_nd.shape
    (m, _) = y_md.shape

    f = ThinPlateSpline(d)
    g = ThinPlateSpline(d)

    src_ctx.reset_tps_params()
    tgt_ctx.reset_tps_params()
    for i, b in enumerate(src_ctx.bend_coefs):
        T = T_vals[i]
        for _ in range(em_iter):
            src_ctx.transform_points()
            tgt_ctx.transform_points()

            xwarped_nd = f.transform_points(x_nd)
            ywarped_md = g.transform_points(y_md)
            gpu_xw = src_ctx.pts_w[test_ind].get()[:n, :]
            gpu_yw = tgt_ctx.pts_w[test_ind].get()[:m, :]
            assert np.allclose(xwarped_nd, gpu_xw, atol=1e-5)
            assert np.allclose(ywarped_md, gpu_yw, atol=1e-5)

            xwarped_nd = gpu_xw
            ywarped_md = gpu_yw

            src_ctx.get_target_points(tgt_ctx, outlierprior, outlierfrac, outliercutoff, T)

            fwddist_nm = ssd.cdist(xwarped_nd, y_md,'euclidean')
            invdist_nm = ssd.cdist(x_nd, ywarped_md,'euclidean')
            prob_nm = outlierprior * np.ones((n+1, m+1), np.float32)
            prob_nm[:n, :m] = np.exp( -(fwddist_nm + invdist_nm) / float(2*T))
            prob_nm[n, m] = outlierfrac * np.sqrt(n * m)

            gpu_corr = src_ctx.corr_rm[test_ind].get()
            gpu_corr = gpu_corr.flatten()
            gpu_corr = gpu_corr[:(n + 1) * (m + 1)].reshape(n+1, m+1).astype(np.float32)

            assert np.allclose(prob_nm[:n, :m], gpu_corr[:n, :m], atol=1e-5)
            prob_nm[:n, :m] = gpu_corr[:n, :m]

            r_coefs = np.ones(n+1, np.float32)
            c_coefs = np.ones(m+1, np.float32)
            a_N = np.ones((n+1),dtype = np.float32)
            a_N[n] = m*outlierfrac
            b_M = np.ones((m+1), dtype = np.float32)
            b_M[m] = n*outlierfrac
            for _ in range(DEFAULT_NORM_ITERS):
                r_coefs = a_N/prob_nm.dot(c_coefs)
                rn_c_coefs = c_coefs
                c_coefs = b_M/r_coefs.dot(prob_nm)
            gpu_r_coefs = src_ctx.r_coefs[test_ind].get()[:n+1].reshape(n+1)
            gpu_c_coefs_cn = src_ctx.c_coefs_cn[test_ind].get()[:m+1].reshape(m+1)
            gpu_c_coefs_rn = src_ctx.c_coefs_rn[test_ind].get()[:m+1].reshape(m+1)
            assert np.allclose(r_coefs, gpu_r_coefs, atol=1e-5)
            assert np.allclose(c_coefs, gpu_c_coefs_cn, atol=1e-5)
            assert np.allclose(rn_c_coefs, gpu_c_coefs_rn, atol=1e-5)

            prob_nm = prob_nm[:n, :m]
            prob_nm *= gpu_r_coefs[:n, None]
            rn_p_nm = prob_nm * gpu_c_coefs_rn[None, :m]
            cn_p_nm = prob_nm * gpu_c_coefs_cn[None, :m]

            wt_n = rn_p_nm.sum(axis=1)
            gpu_corr_cm = src_ctx.corr_cm[test_ind].get().flatten()[:(n+1)*(m+1)]
            gpu_corr_cm = gpu_corr_cm.reshape(m+1, n+1)## b/c it is column major
            assert np.allclose(wt_n, gpu_corr_cm[m, :n], atol=1e-4)

            inlier = wt_n > outliercutoff
            xtarg_nd = np.empty((n, DATA_DIM), np.float32)
            xtarg_nd[inlier, :] = rn_p_nm.dot(y_md)[inlier, :]
            xtarg_nd[~inlier, :] = xwarped_nd[~inlier, :]

            wt_m = cn_p_nm.sum(axis=0)
            assert np.allclose(wt_m, gpu_corr[n, :m], atol=1e-4)

            inlier = wt_m > outliercutoff
            ytarg_md = np.empty((m, DATA_DIM), np.float32)
            ytarg_md[inlier, :] = cn_p_nm.T.dot(x_nd)[inlier, :]
            ytarg_md[~inlier, :] = ywarped_md[~inlier, :]

            xt_gpu = src_ctx.pts_t[test_ind].get()[:n, :]
            yt_gpu = tgt_ctx.pts_t[test_ind].get()[:m, :]
            assert np.allclose(xtarg_nd, xt_gpu, atol=1e-4)
            assert np.allclose(ytarg_md, yt_gpu, atol=1e-4)

            src_ctx.update_transform(b)
            tgt_ctx.update_transform(b)

            f_p_mat = src_ctx.proj_mats[b][test_ind].get()[:n+d+1, :n]
            f_o_mat = src_ctx.offset_mats[b][test_ind].get()[:n+d+1]
            b_p_mat = tgt_ctx.proj_mats[b][0].get()[:m+d+1, :m]
            b_o_mat = tgt_ctx.offset_mats[b][0].get()[:m+d+1]
            f_params = f_p_mat.dot(xtarg_nd) + f_o_mat
            g_params = b_p_mat.dot(ytarg_md) + b_o_mat

            gpu_fparams = src_ctx.tps_params[test_ind].get()[:n+d+1]
            gpu_gparams = tgt_ctx.tps_params[test_ind].get()[:m+d+1]
            assert np.allclose(f_params, gpu_fparams, atol=1e-4)
            assert np.allclose(g_params, gpu_gparams, atol=1e-4)

            set_ThinPlateSpline(f, x_nd, gpu_fparams)
            set_ThinPlateSpline(g, y_md, gpu_gparams)

    f._cost = tps.tps_cost(f.lin_ag, f.trans_g, f.w_ng, f.x_na, xtarg_nd, 1)
    g._cost = tps.tps_cost(g.lin_ag, g.trans_g, g.w_ng, g.x_na, ytarg_md, 1)

    gpu_cost = src_ctx.bidir_tps_cost(tgt_ctx)
    cpu_cost = f._cost + g._cost
    assert np.isclose(gpu_cost[test_ind], cpu_cost, atol=1e-4)

def parse_arguments():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_file", type=str, default='../data/actions.h5')
    parser.add_argument("--sync", action='store_true')
    parser.add_argument("--n_copies", type=int, default=1)
    parser.add_argument("--test", action='store_true')
    parser.add_argument("--test_full", action='store_true')
    parser.add_argument("--timing_runs", type=int, default=100)
    return parser.parse_args()

if __name__=='__main__':
    reset_cuda()
    args = parse_arguments()
    Globals.sync = args.sync
    src_ctx = GPUContext()
    for _ in range(args.n_copies):
        src_ctx.read_h5(args.input_file)
    f = h5py.File(args.input_file, 'r')
    tgt_cld = downsample_cloud(f['demo1-seg00']['cloud_xyz'][:])
    f.close()
    scaled_tgt_cld, _ = unit_boxify(tgt_cld)
    tgt_ctx = TgtContext(src_ctx)
    tgt_ctx.set_cld(scaled_tgt_cld)
    if args.test_full:
        src_ctx.unit_test(tgt_ctx)
        print "unit tests passed, doing full check on batch tps rpm"
        for i in range(src_ctx.N):
            sys.stdout.write("\rtesting source cloud {}".format(i))
            sys.stdout.flush()
            test_batch_tps_rpm_bij(src_ctx, tgt_ctx, test_ind=i)
        print""
        print "tests succeeded!"
        sys.exit()
    if args.test:
        src_ctx.unit_test(tgt_ctx)
        print "testing batch tps_rps"
        test_batch_tps_rpm_bij(src_ctx, tgt_ctx)
        print "test succeeded!!"
        sys.exit()
    times = []
    print "batchtps initialized"
    for i in range(args.timing_runs):
        sys.stdout.write("\rRunning Timing test {}/{}".format(i, args.timing_runs))
        sys.stdout.flush()
        start = time.time()
        tgt_ctx.set_cld(scaled_tgt_cld)
        c = batch_tps_rpm_bij(src_ctx, tgt_ctx)
        time_taken = time.time() - start
        times.append(time_taken)
    print "\nTiming Tests Complete"
    print "Batch Size:\t\t\t", src_ctx.N
    print "Mean Compute Time per Batch:\t", np.mean(times)
    print "BiDirectional TPS fits/second:\t", float(args.timing_runs * src_ctx.N) / np.sum(times)

import numpy,import scipy,import pycuda,from pycuda


#!/usr/bin/env python3
import sys, getopt, math, os, time, traceback
import numpy as np

################################# Arguments ###################################
# Default parameters
n_traj = 100
num_rep = 10
mutant_type_list = ['fast', 'slow']
n_window = 200
co_msm_data_file = ''
post_msm_data_file = ''

# read control file
ctrlfile = ''

if len(sys.argv) == 1:
    print(usage)
    sys.exit()

try:
    opts, args = getopt.getopt(sys.argv[1:],"hf:", ["ctrlfile="])
except getopt.GetoptError:
    print(usage)
    sys.exit()
for opt, arg in opts:
    if opt == '-h':
        print(usage)
        sys.exit()
    elif opt in ("-f", "--ctrlfile"):
        ctrlfile = arg

if not os.path.exists(ctrlfile):
    print('Error: cannot find control file ' + ctrlfile + '.')
    sys.exit()

file_object = open(ctrlfile,'r')
try:
    for line in file_object:
        line = line.strip()
        if not line:
            # This is a blank line
            continue
        if line.startswith('#'):
            # This is a comment line
            continue
        if line.startswith('n_traj'):
            words = line.split('=')
            n_traj = int(words[1].strip())
            continue
        if line.startswith('num_rep'):
            words = line.split('=')
            num_rep = int(words[1].strip())
            continue
        if line.startswith('mutant_type_list'):
            words = line.split('=')
            mutant_type_list = words[1].strip().split()
            continue
        if line.startswith('n_window'):
            words = line.split('=')
            n_window = int(words[1].strip())
            continue
        if line.startswith('co_msm_data_file'):
            words = line.split('=')
            co_msm_data_file = words[1].strip()
            continue
        if line.startswith('post_msm_data_file'):
            words = line.split('=')
            post_msm_data_file = words[1].strip()
            continue
finally:
     file_object.close()

################################# Functions ###################################
def get_pathways(meta_dtrajs, start_states, end_states):
    pathways = {}
    states_on_pathway = []
    start_states = [str(s+1) for s in start_states]
    end_states = [str(s+1) for s in end_states]

    for traj_idx, md in enumerate(meta_dtrajs):
        path = []
        for idx, mdi in enumerate(md):
            if str(mdi+1) in start_states:
                path.append(str(mdi+1))
                break
        if idx == len(md)-1:
            continue

        for mdi in md[idx+1:]:
            tag_find = False
            for pi in range(len(path)):
                if path[pi] == str(mdi+1):
                    path = path[0:pi+1]
                    tag_find = True
                    break
            if not tag_find:
                path.append(str(mdi+1))

        if path[-1] not in end_states:
            continue

        for p in path:
            if not int(p) in states_on_pathway:
                states_on_pathway.append(int(p))

        path = ' -> '.join(path)
        if not path in pathways.keys():
            pathways[path] = 1
        else:
            pathways[path] += 1

    tot_num = 0
    for path in pathways.keys():
        tot_num += pathways[path]
    for path in pathways.keys():
        pathways[path] /= tot_num

    sort_pathways = sorted(pathways.items(), key=lambda x: x[1], reverse=True)

    states_on_pathway.sort()
    return [sort_pathways, states_on_pathway]

################################## MAIN #######################################
# Read co-trans metastable dtrajs
co_meta_dtrajs = np.load(co_msm_data_file, allow_pickle=True)['meta_dtrajs']
co_n_states = 0
for i in range(len(co_meta_dtrajs)):
    if np.max(co_meta_dtrajs[i]) > co_n_states:
        co_n_states = np.max(co_meta_dtrajs[i])
co_n_states += 1
# Read post-trans metastable dtrajs
post_meta_dtrajs = np.load(post_msm_data_file, allow_pickle=True)['meta_dtrajs']
max_T_len = 0
post_n_states = 0
for i in range(len(post_meta_dtrajs)):
    if post_meta_dtrajs[i].shape[0] > max_T_len:
        max_T_len = post_meta_dtrajs[i].shape[0]
    if np.max(post_meta_dtrajs[i]) > post_n_states:
        post_n_states = np.max(post_meta_dtrajs[i])
post_n_states += 1
post_meta_dtrajs_extended = []
for i in range(len(post_meta_dtrajs)):
    (N, be) = np.histogram(post_meta_dtrajs[i][-n_window:], bins=np.arange(-0.5, post_n_states, 1))
    meta_dtraj_last = np.argwhere(N == np.max(N))[0][0]
    mde = []
    for j in range(max_T_len):
        if j >= len(post_meta_dtrajs[i]):
            state_0 = meta_dtraj_last
        else:
            state_0 = post_meta_dtrajs[i][j]
        mde.append(state_0)
    post_meta_dtrajs_extended.append(mde)

print('Total number of states: %d'%(co_n_states + post_n_states))

# combine dtrajs
meta_dtrajs = []
for i in range(len(co_meta_dtrajs)):
    for j in range(num_rep):
        md = np.hstack((co_meta_dtrajs[i], post_meta_dtrajs_extended[i*num_rep+j] + co_n_states))
        meta_dtrajs.append(md)
meta_dtrajs = np.array(meta_dtrajs, dtype=object)

mtype2trajid = [np.arange(n_traj*num_rep*i_ax, n_traj*num_rep*(i_ax+1)).astype(int) for i_ax, mutant_type in enumerate(mutant_type_list)]

# analysis MSM for each mutant
fo = open('pathways.dat', 'w')
for i_ax, mutant_type in enumerate(mutant_type_list):
    # flux analysis
    A = [0]
    B = [md[-1] for md in meta_dtrajs[mtype2trajid[i_ax]]]
    B = np.unique(B)

    [pathways, state_on_pathway] = get_pathways(meta_dtrajs[mtype2trajid[i_ax]], A, B)

    fo.write('%s pathways:\n'%mutant_type)
    fo.write('%-12s %s\n'%('percentage','path'))
    fo.write('-------------------------------------\n')
    for path in pathways:
        fo.write('%-12s %-s\n'%('%.2f %%'%(path[1] * 100), path[0]))
    fo.write('\n')


# -*- coding: utf-8 -*-
"""
Created on Tue Jun 29 09:38:12 2021

@author: ruy
"""
from __future__ import annotations
from abc import ABC
from dataclasses import dataclass, field
from enum import Enum
from functools import cached_property as property
import numpy as np
from dataclass_tools.tools import (
    DESERIALIZER_OPTIONS,
    DeSerializerOptions,
    NamePrint,
    PrintMetadata,
    serialize_dataclass,
)

from pylatex import NoEscape, Quantity

from .abrevitation_registry import abv_registry
from .common_field_options import NAME_OPTIONS
from .constants import GRAVITY

class TypeOfService(str, Enum):
    """Vessels service type in accordance to GL HSC 2012 rules."""

    PASSENGER = "PASSENGER"
    FERRY = "FERRY"
    CARGO = "CARGO"
    SUPPLY = "SUPPLY"
    PILOT = "PILOT"
    RESCUE = "RESCUE"

class ServiceRange(str, Enum):
    """
    Vessels service range in accordance to GL HSC 2012 rules.
    USR - Unrestricted service range
    RSA_200 - 200 nm range
    RSA_50 - 50 nm range
    RSA_20 - 20 nm range
    RSA_SW - sheltered waters
    """

    USR = "USR"
    RSA_200 = "RSA (200)"
    RSA_50 = "RSA (50)"
    RSA_20 = "RSA (20)"
    RSA_SW = "RSA (SW)"

SPEED_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(long_name="Speed", abreviation="S", units="kt")
)
abv_registry.append(SPEED_OPTIONS)
DISPLACEMENT_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Displacement", abreviation=NoEscape(r"$\Delta$"), units="t"
    )
)
abv_registry.append(DISPLACEMENT_OPTIONS)
LENGTH_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(long_name="Length", abreviation="L", units="m")
)
abv_registry.append(LENGTH_OPTIONS)
BEAM_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(long_name="Beam", abreviation="B", units="m")
)
abv_registry.append(BEAM_OPTIONS)
FWD_PERP_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Foward Perpendicular",
        abreviation=NoEscape(r"F\textsubscript{P}"),
        units="m",
    )
)
abv_registry.append(FWD_PERP_OPTIONS)
AFT_PERP_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Aft Perpendicular",
        abreviation=NoEscape(r"A\textsubscript{P}"),
        units="m",
    )
)
abv_registry.append(AFT_PERP_OPTIONS)
DRAFT_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Draft",
        abreviation="H",
        units="m",
    )
)
abv_registry.append(DRAFT_OPTIONS)
Z_BASELINE_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Z from base line",
        abreviation=NoEscape(r"z\textsubscript{BL}"),
        units="m",
    )
)
abv_registry.append(Z_BASELINE_OPTIONS)
BLOCK_COEF_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Block coefficient",
        abreviation=NoEscape(r"C\textsubscript{B}"),
        units="",
    )
)
abv_registry.append(BLOCK_COEF_OPTIONS)
WATER_PLANE_AREA_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Water plane area",
        abreviation=NoEscape(r"WP\textsubscript{A}"),
        units="m**2",
    )
)
abv_registry.append(WATER_PLANE_AREA_OPTIONS)
LCG_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Longitudinal center of gravity",
        abreviation=NoEscape(r"L\textsubscript{cg}"),
        units="m",
    )
)
abv_registry.append(LCG_OPTIONS)
DEADRISE_LCG_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Deadrise at longitudinal center of gravity",
        abreviation=NoEscape(r"$\alpha$\textsubscript{d}"),
        units="degree",
    )
)
abv_registry.append(DEADRISE_LCG_OPTIONS)
DIST_HULL_CL_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Distance between hulls centerline",
        abreviation=NoEscape(r"B\textsubscript{CL}"),
        units="m",
    )
)
abv_registry.append(DIST_HULL_CL_OPTIONS)
TYPE_OF_SERVICE_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Type of service",
        abreviation="ToS",
    )
)
abv_registry.append(TYPE_OF_SERVICE_OPTIONS)
SERVICE_RANGE_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Service range",
        abreviation="SR",
    )
)
abv_registry.append(SERVICE_RANGE_OPTIONS)
VERT_ACG_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Center of gravity's vertical acceleration",
        abreviation=NoEscape(r"a\textsubscript{CG}"),
        units="g",
    )
)
abv_registry.append(VERT_ACG_OPTIONS)
SHEAR_FORCE_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Transverse shear force",
        abreviation=NoEscape(r"T\textsubscript{bt}"),
        units="kN",
    )
)
abv_registry.append(SHEAR_FORCE_OPTIONS)
TRANSVERSE_BENDING_MOMENT_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Transverse bending moment",
        abreviation=NoEscape(r"M\textsubscript{bt}"),
        units="kN*m",
    )
)
abv_registry.append(TRANSVERSE_BENDING_MOMENT_OPTIONS)
TRANSVERSE_TORSIONAL_MOMENT_OPTIONS = DeSerializerOptions(
    metadata=PrintMetadata(
        long_name="Transverse torsional moment",
        abreviation=NoEscape(r"M\textsubscript{tt}"),
        units="kN*m",
    )
)
abv_registry.append(TRANSVERSE_TORSIONAL_MOMENT_OPTIONS)

@dataclass
class VesselLoads:
    vert_acg: float = field(metadata={DESERIALIZER_OPTIONS: VERT_ACG_OPTIONS})
    shear_force: float = field(metadata={DESERIALIZER_OPTIONS: SHEAR_FORCE_OPTIONS})
    trans_bend_moment: float = field(
        metadata={DESERIALIZER_OPTIONS: TRANSVERSE_BENDING_MOMENT_OPTIONS}
    )
    trans_tors_moment: float = field(
        metadata={DESERIALIZER_OPTIONS: TRANSVERSE_TORSIONAL_MOMENT_OPTIONS}
    )

@dataclass
class Monohull:
    """Vessel for the 2012 German Lloyds High Speed Craft
    scantling rules.
    """

    name: str = field(metadata={DESERIALIZER_OPTIONS: NAME_OPTIONS})
    speed: float = field(metadata={DESERIALIZER_OPTIONS: SPEED_OPTIONS})
    displacement: float = field(metadata={DESERIALIZER_OPTIONS: DISPLACEMENT_OPTIONS})
    length: float = field(metadata={DESERIALIZER_OPTIONS: LENGTH_OPTIONS})
    beam: float = field(metadata={DESERIALIZER_OPTIONS: BEAM_OPTIONS})
    fwd_perp: float = field(metadata={DESERIALIZER_OPTIONS: FWD_PERP_OPTIONS})
    aft_perp: float = field(metadata={DESERIALIZER_OPTIONS: AFT_PERP_OPTIONS})
    draft: float = field(metadata={DESERIALIZER_OPTIONS: DRAFT_OPTIONS})
    z_baseline: float = field(metadata={DESERIALIZER_OPTIONS: Z_BASELINE_OPTIONS})
    block_coef: float = field(metadata={DESERIALIZER_OPTIONS: BLOCK_COEF_OPTIONS})
    water_plane_area: float = field(
        metadata={DESERIALIZER_OPTIONS: WATER_PLANE_AREA_OPTIONS}
    )
    lcg: float = field(metadata={DESERIALIZER_OPTIONS: LCG_OPTIONS})
    deadrise_lcg: float = field(metadata={DESERIALIZER_OPTIONS: DEADRISE_LCG_OPTIONS})

    type_of_service: TypeOfService = field(
        metadata={DESERIALIZER_OPTIONS: TYPE_OF_SERVICE_OPTIONS},
        default=TypeOfService.PASSENGER,
    )
    service_range: ServiceRange = field(
        metadata={DESERIALIZER_OPTIONS: SERVICE_RANGE_OPTIONS},
        default=ServiceRange.USR,
    )

    # Table C3.3.1
    @property
    def serv_type_coef(self):
        table = {
            TypeOfService.PASSENGER: 0.24,
            TypeOfService.FERRY: 0.24,
            TypeOfService.CARGO: 0.24,
            TypeOfService.SUPPLY: 0.36,
            TypeOfService.PILOT: 0.5,
            TypeOfService.RESCUE: 0.6,
        }
        return table[self.type_of_service]

    # Table C3.3.1
    @property
    def serv_range_coef(self):
        table = {
            ServiceRange.USR: 1.0,
            ServiceRange.RSA_200: 0.9,
            ServiceRange.RSA_50: 0.75,
            ServiceRange.RSA_20: 0.66,
            ServiceRange.RSA_SW: 0.6,
        }
        return table[self.service_range]

    @property
    def z_waterline(self):
        return self.z_baseline + self.draft

    @property
    def midship(self):
        return (self.fwd_perp + self.aft_perp) / 2

    @property
    def x_pos_cg(self):
        return (self.lcg - self.aft_perp) / (self.fwd_perp - self.aft_perp)

    @property
    def sp_len_ratio(self) -> float:
        return self.speed / self.length**0.5

    # C3.3.1
    @property
    def vert_acg(self):
        """Vertical acceletation at LCGm, output in g (9.81 m/s2)"""
        acg = self.serv_type_coef * self.serv_range_coef * self.sp_len_ratio
        if self.type_of_service == TypeOfService.PASSENGER:
            return min([1.0, acg])
        return acg

    @property
    def max_wave_height(self):
        """C3.3.3 Assessment of limit operating conditions"""
        return (
            5
            * np.max([self.vert_acg, 1])
            / self.speed
            * self.length**1.5
            / (6 + 0.14 * self.length)
        )

    @property
    def sig_wave_height(self):
        """C3.3.3.2 Limitation imposed by vertical acceleration at LCG"""
        return 10.9 * self.vert_acg * self.coef_kcat * self.coef_kh / self.coef_kf**2

    # C3.3.3.2
    @property
    def coef_kcat(self):
        return 1.0

    # C3.3.3.2
    @property
    def coef_kf(self):
        return 3.23 / self.length * (2.43 * self.length**0.5 + self.speed)

    # C3.3.3.2
    @property
    def coef_kt(self):
        return (
            4.6 * self.water_plane_area / self.displacement * (self.x_pos_cg) ** 0.5
        ) ** 0.5

    # C3.3.3.2
    @property
    def coef_k(self):
        return self.coef_kf / self.coef_kt

    # C3.3.3.2
    @property
    def coef_kh(self):
        return self.coef_k**0.35 * ((1 / self.coef_k**2 - 0.11) ** 2 + 1) ** 0.5

@dataclass
class Catamaran(Monohull):
    dist_hull_cl: float = field(
        metadata={DESERIALIZER_OPTIONS: DIST_HULL_CL_OPTIONS}, default=0
    )

    # C3.3.3.2
    @property
    def coef_kcat(self):
        return np.max(
            [1 + (self.dist_hull_cl - self.max_wave_height) / self.length, 1.0]
        )

    # C3.4.2.3
    @property
    def transverse_bending_moment(self):
        """Transverse bending moment in kN*m."""
        return self.displacement * self.dist_hull_cl * self.vert_acg * GRAVITY / 5

    # C3.4.2.3
    @property
    def transverse_shear_force(self):
        """Transverse shear force in kN."""
        return self.displacement * self.vert_acg * GRAVITY / 4

    # C3.4.2.4
    @property
    def transverse_torsional_moment(self):
        """Transverse torsional moment in kN*m."""
        return 0.125 * self.displacement * self.length * self.vert_acg * GRAVITY

    @property
    def loads_names(self):
        return [
            "vert_acg",
            "shear_force",
            "transverse_bending_moment",
            "transverse_torsional_moment",
        ]

    @property
    def loads_asdict(self):
        loads = VesselLoads(
            vert_acg=self.vert_acg,
            shear_force=self.transverse_shear_force,
            trans_bend_moment=self.transverse_bending_moment,
            trans_tors_moment=self.transverse_torsional_moment,
        )
        return serialize_dataclass(loads, printing_format=True, include_names=True)