"""
Inference with different ADC precisions
"""
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from .utee import wage_quantizer
from .quant_modules import WQ, AQ, stats_quant

def decimal2binary(weight_q, bitWeight, cellBit):
    cellRange = 2**cellBit
    remainder_list = torch.Tensor([]).type_as(weight_q)

    for k in range(int(bitWeight/cellBit)):
        remainder = torch.fmod(weight_q, cellRange)
        remainder = remainder.unsqueeze(0)
        remainder_list = torch.cat((remainder_list, remainder), dim=0)
        weight_q = torch.round((weight_q-remainder.squeeze(0))/cellRange)
    return remainder_list

def bit2cond(bitWeight, hrs, lrs):
    """
    Draft: replace the binary values to conductance measurement
    """
    level0 = torch.ones(bitWeight[bitWeight==0].size()).mul(hrs)
    level1 = torch.ones(bitWeight[bitWeight==1].size()).mul(lrs)

    bitWeight[bitWeight==0] = level0.cuda()
    bitWeight[bitWeight==1] = level1.cuda()

    bitWeight = bitWeight.clamp(0)
    return bitWeight

def program_noise_cond(weight_q, weight_b, hrs, lrs, sensitive_lv):
    wb = torch.zeros_like(weight_b)
    weight_cond = bit2cond(weight_b, hrs, lrs)  # typical values

    for ii in range(len(weight_q.unique())):

        if len(weight_q.unique()) == 1:
            ii = int(weight_q.unique().item())

        idx_4b = weight_q.eq(ii)
        wb_ii = weight_cond[:, idx_4b]
        wbin_ii = weight_b[:, idx_4b]

        # noises
        noise = np.load(f"/home/mengjian/Desktop/ASU_research/SWIPE_analysis/prob/SWIPE/noSWIPE_25Times_raw/level{ii}_raw.npy") # 1.66
        # swipe = np.load(f"/home/mengjian/Desktop/ASU_research/SWIPE_analysis/prob/SWIPE/Level_4x16_SWIPE_250nPW_chip14_raw_in_16lvl/level{ii}_raw.npy")

        # rescale the programming noise with new lrs （1.66 to 1.11)
        # noise = noise * (1.11e-4/1.66e-4)
        # swipe = swipe * (1.11e-4/1.66e-4)

        # noise = np.load(f"/home/mengjian/Desktop/ASU_research/SWIPE_analysis/prob/SWIPE/DNN_SWIPE_101021/Level_4x16_noSWIPE_250nPW_chip14_raw_in_16lvl_10080928/level{ii}_raw.npy")
        # swipe = np.load(f"/home/mengjian/Desktop/ASU_research/SWIPE_analysis/prob/SWIPE/DNN_SWIPE_102021/Level_4x16_SWIPE_250nPW_chip14_raw_in_16lvl_10201845/level{ii}_raw.npy")
        # swipe = np.load(f"/home/mengjian/Desktop/ASU_research/SWIPE_analysis/prob/SWIPE/DNN_SWIPE_102021/09191845/level{ii}_raw.npy")   # 1.11
        swipe = np.load(f"/home/mengjian/Desktop/ASU_research/SWIPE_analysis/prob/SWIPE/DNN_SWIPE_102021/Level_4x16_SWIPE_250nPW_chip14_raw_in_16lvl_10201845/level{ii}_raw.npy")   # 1.11
        swipe = swipe * (1.66e-4/1.11e-4)

        # sizes
        _, numel = wb_ii.size()

        bit_idx = np.arange(noise.shape[0])
        random_idx = np.random.choice(bit_idx, size=(numel))

        bit_random_noise = noise[random_idx, :].T
        swipe_random_noise = swipe[random_idx, :].T

        wb_cond = torch.from_numpy(bit_random_noise).float()
        swipe_cond = torch.from_numpy(swipe_random_noise).float()

        wb_cond = torch.flip(wb_cond, dims=[0])
        swipe_cond = torch.flip(swipe_cond, dims=[0])

        if not ii in sensitive_lv:
            wb[:, idx_4b] = swipe_cond.cuda()   # SWIPE scheme
        else:
            wb[:, idx_4b] = wb_cond.cuda()      # Non SWIPE scheme

        # # statistics
        # hrs = wb_cond[wb_ii == 0]
        # lrs = wb_cond[wb_ii == 1]
        # print("\nLevel {}; Average HRS={}; Average LRS={}; dummy={}".format(ii, hrs.mean(), lrs.mean(), dummy))
    return wb

class RRAMConv2d(nn.Conv2d):
    r"""
    NeuroSim-based RRAM inference with low precision weights and activations
    """
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True,
                wl_input=8, wl_weight=8, subArray=128, inference=0, cellBit=1, ADCprecision=5, sensitive_lv=0):
        super(RRAMConv2d, self).__init__(in_channels, out_channels, kernel_size,
                                      stride, padding, dilation, groups, bias)
        self.wl_input = wl_input
        self.inference = inference
        self.wl_weight = wl_weight
        self.cellBit = cellBit
        self.ADCprecision = ADCprecision
        self.layer_idx = 0
        self.iter = 0
        self.init = True
        self.subArray = subArray

        # quantization
        self.wbit = wl_weight
        self.abit = wl_input
        self.weight_quant = WQ(wbit=wl_weight)
        self.act_quant = AQ(abit=wl_input, act_alpha=torch.tensor(10.0))

        # conductance
        self.hrs = 1e-6
        # self.lrs = 1.11e-04
        self.lrs = 1.66e-04
        self.nonideal_unit = self.lrs - self.hrs
        self.sensitive_lv = sensitive_lv

    def _act_quant(self, input):
        act_alpha = self.act_quant.act_alpha
        input = torch.where(input < act_alpha, input, act_alpha)

        with torch.no_grad():
            scale = (2**self.abit - 1) / act_alpha

        input_div = input.mul(scale)
        input_q = input_div.round()
        return input_q, scale

    def forward(self, input: Tensor) -> Tensor:
        # quantization
        wq, w_scale = stats_quant(self.weight.data, nbit=self.wbit, dequantize=False)
        wq = wq.add(2 ** (self.wbit - 1) - 1)
        wd = torch.ones_like(wq).mul(2 ** (self.wbit - 1) - 1)

        # decomposition
        wqb_list = decimal2binary(wq, bitWeight=self.wbit, cellBit=self.cellBit)
        wdb_list = decimal2binary(wd, bitWeight=self.wbit, cellBit=self.cellBit)

        wqb_list = program_noise_cond(wq, wqb_list, hrs=self.hrs, lrs=self.lrs, sensitive_lv=self.sensitive_lv)
        wdb_list = program_noise_cond(wd, wdb_list, hrs=self.hrs, lrs=self.lrs, sensitive_lv=self.sensitive_lv)

        # input quantization
        xq, x_scale = self._act_quant(input)
        cellRange = 2**self.cellBit

        # targeted output size
        odim = math.floor((xq.size(2) + 2*self.padding[0] - self.dilation[0] * (wq.size(2)-1)-1)/self.stride[0] + 1)
        output = torch.zeros((xq.size(0), wq.size(0), odim, odim)).cuda()
        for i in range(wq.size(2)):
            for j in range(wq.size(3)):
                numSubArray = wq.shape[1] // self.subArray

                if numSubArray == 0:
                    mask = torch.zeros_like(wq)
                    mask[:,:,i,j] = 1
                    xq, x_scale = self._act_quant(input)
                    outputIN = torch.zeros_like(output)
                    xb_list = []
                    for z in range(int(self.abit)):
                        xb = torch.fmod(xq, 2)
                        xq = torch.round((xq-xb)/2)
                        macs = torch.zeros_like(output).cuda()

                        xb_list.append(xb)
                        for k in range(int(self.wbit/self.cellBit)):
                            wqb = wqb_list[k]
                            wdb = wdb_list[k]

                            outputPartial = F.conv2d(xb, wqb*mask, self.bias, self.stride, self.padding, self.dilation, self.groups)
                            outputOffset = F.conv2d(xb, wdb*mask, self.bias, self.stride, self.padding, self.dilation, self.groups)
                            scaler = cellRange**k

                            maci = outputPartial - outputOffset
                            maci = maci.div(self.nonideal_unit)

                            macs = macs + maci * scaler

                        scalerIN = 2**z
                        outputIN = outputIN + macs * scalerIN
                    output = output + outputIN / x_scale
                else:
                    xq, x_scale = self._act_quant(input)
                    outputIN = torch.zeros_like(output)
                    for z in range(int(self.abit)):
                        xb = torch.fmod(xq, 2)
                        xq = torch.round((xq-xb)/2)
                        total_macs = torch.zeros_like(output)

                        for s in range(numSubArray):
                            mask = torch.zeros_like(wq)
                            mask[:,(s*self.subArray):(s+1)*self.subArray, i, j] = 1
                            macs = torch.zeros_like(output).cuda()

                            for k in range(int(self.wbit/self.cellBit)):
                                wqb = wqb_list[k]
                                wdb = wdb_list[k]

                                outputPartial = F.conv2d(xb, wqb*mask, self.bias, self.stride, self.padding, self.dilation, self.groups)
                                outputOffset = F.conv2d(xb, wdb*mask, self.bias, self.stride, self.padding, self.dilation, self.groups)
                                scaler = cellRange**k

                                maci = outputPartial - outputOffset
                                # ADC
                                maci = wage_quantizer.LinearQuantizeOut(maci, bit=self.ADCprecision, lb=maci.min(), ub=maci.max())
                                maci = maci.div(self.nonideal_unit)
                                macs = macs + maci * scaler

                            total_macs = total_macs.add(macs)
                        scalerIN = 2**z
                        outputIN = outputIN + total_macs * scalerIN
                    output = output + outputIN / x_scale
        output = output / w_scale
        return output


import logging
import multiprocessing
import sys
import traceback
from multiprocessing import Pool

import numpy as np
import ellc
import matplotlib.pyplot as plt
from lcbuilder.lcbuilder_class import LcBuilder
from lcbuilder.objectinfo.InputObjectInfo import InputObjectInfo
from lcbuilder.objectinfo.MissionInputObjectInfo import MissionInputObjectInfo
from lcbuilder.objectinfo.preparer.MissionLightcurveBuilder import MissionLightcurveBuilder
from lcbuilder.objectinfo.MissionObjectInfo import MissionObjectInfo
import wotan
from matplotlib.ticker import FormatStrFormatter
from foldedleastsquares import transitleastsquares, DefaultTransitTemplateGenerator
from foldedleastsquares import transit_mask, cleaned_array
import astropy.constants as ac
import astropy.units as u
import lightkurve as lk
import os
import re
import pandas as pd

from tkmatrix.inject_model import InjectModel

class MATRIX:
    """
    MATRIX: Multi-phAse Transits Recovery from Injected eXoplanets
    """
    object_info = None
    SDE_ROCHE = 2000
    lcbuilder = LcBuilder()
    MIN_SEARCH_PERIOD = 0.5

    def __init__(self, target, sectors, dir, preserve=False, star_info=None, file=None, exposure_time=None,
                 initial_mask=None, initial_transit_mask=None,
                 eleanor_corr_flux='pca_flux', outliers_sigma=None, high_rms_enabled=True, high_rms_threshold=2.5,
                 high_rms_bin_hours=4, smooth_enabled=False,
                 auto_detrend_enabled=False, auto_detrend_method="cosine", auto_detrend_ratio=0.25,
                 auto_detrend_period=None, prepare_algorithm=None, cache_dir=os.path.expanduser('~') + "/",
                 oscillation_reduction=False, oscillation_min_snr=4, oscillation_amplitude_threshold=0.001,
                 oscillation_ws_percent=0.01, oscillation_min_period=0.002, oscillation_max_period=0.2,
                 cores=multiprocessing.cpu_count() - 1
                 ):
        assert target is not None and isinstance(target, str)
        assert sectors is not None and (sectors == 'all' or isinstance(sectors, list))
        assert exposure_time is not None and isinstance(exposure_time, (int, float))
        assert initial_transit_mask is None or isinstance(initial_transit_mask, list)
        self.id = target
        self.dir = dir
        self.sectors = sectors
        self.star_info = star_info
        self.exposure_time = exposure_time
        self.preserve = preserve
        self.file = file
        self.eleanor_corr_flux = eleanor_corr_flux
        self.initial_mask = initial_mask
        self.initial_transit_mask = initial_transit_mask
        self.star_info = star_info
        self.outliers_sigma = outliers_sigma
        self.high_rms_enabled = high_rms_enabled
        self.high_rms_threshold = high_rms_threshold
        self.high_rms_bin_hours = high_rms_bin_hours
        self.smooth_enabled = smooth_enabled
        self.auto_detrend_enabled = auto_detrend_enabled
        self.auto_detrend_method = auto_detrend_method
        self.auto_detrend_ratio = auto_detrend_ratio
        self.auto_detrend_period = auto_detrend_period
        self.oscillation_reduction = oscillation_reduction
        self.oscillation_min_snr = oscillation_min_snr
        self.oscillation_amplitude_threshold = oscillation_amplitude_threshold
        self.oscillation_ws_percent = oscillation_ws_percent
        self.oscillation_min_period = oscillation_min_period
        self.oscillation_max_period = oscillation_max_period
        self.prepare_algorithm = prepare_algorithm
        self.cache_dir = cache_dir
        self.cores = cores

    def retrieve_object_data(self, inject_dir=None):
        self.object_info = self.lcbuilder.build_object_info(self.id, None, self.sectors, self.file, self.exposure_time,
                                                       None, None,
                                                       self.star_info, None,
                                                       self.eleanor_corr_flux, self.outliers_sigma,
                                                       False, self.high_rms_threshold,
                                                       self.high_rms_bin_hours, False,
                                                       False, self.auto_detrend_method,
                                                       self.auto_detrend_ratio, self.auto_detrend_period,
                                                       self.prepare_algorithm, False,
                                                       self.oscillation_min_snr, self.oscillation_amplitude_threshold,
                                                       self.oscillation_ws_percent, self.oscillation_min_period,
                                                       self.oscillation_max_period)
        if inject_dir is None:
            inject_dir = self.build_inject_dir()
        self.lc_build = self.lcbuilder.build(self.object_info, inject_dir, self.cache_dir)
        if self.star_info is None:
            self.star_info = self.lc_build.star_info
        self.ab = self.star_info.ld_coefficients
        self.mass = self.star_info.mass
        self.massmin = self.star_info.mass_min
        self.massmax = self.star_info.mass_max
        self.radius = self.star_info.radius
        self.radiusmin = self.star_info.radius_min
        self.radiusmax = self.star_info.radius_max
        # units for ellc
        self.rstar = self.star_info.radius * u.R_sun
        self.mstar = self.star_info.mass * u.M_sun
        self.mstar_min = self.star_info.mass_min * u.M_sun
        self.mstar_max = self.star_info.mass_max * u.M_sun
        self.rstar_min = self.star_info.radius_min * u.R_sun
        self.rstar_max = self.star_info.radius_max * u.R_sun
        return inject_dir

    def retrieve_object_data_for_recovery(self, inject_dir, recovery_file):
        self.__setup_logging(inject_dir)
        self.object_info = self.lcbuilder.build_object_info("", None, None, recovery_file, self.exposure_time,
                                                       self.initial_mask, self.initial_transit_mask,
                                                       self.star_info, None,
                                                       self.eleanor_corr_flux, self.outliers_sigma,
                                                       self.high_rms_enabled, self.high_rms_threshold,
                                                       self.high_rms_bin_hours, self.smooth_enabled,
                                                       self.auto_detrend_enabled, self.auto_detrend_method,
                                                       self.auto_detrend_ratio, self.auto_detrend_period,
                                                       self.prepare_algorithm, self.oscillation_reduction,
                                                       self.oscillation_min_snr, self.oscillation_amplitude_threshold,
                                                       self.oscillation_ws_percent, self.oscillation_min_period,
                                                       self.oscillation_max_period)

        if self.object_info.reduce_simple_oscillations and \
                self.object_info.oscillation_max_period < self.object_info.oscillation_min_period:
            logging.info("Stellar oscillation period has been set to empty. Defaulting to 1/3 the minimum search period")
            self.object_info.oscillation_max_period = self.MIN_SEARCH_PERIOD / 3
        self.lc_build = self.lcbuilder.build(self.object_info, inject_dir, self.cache_dir)

    def build_inject_dir(self):
        inject_dir = self.dir + "/" + self.object_info.mission_id().replace(" ", "") + "_ir/"
        index = 0
        while os.path.exists(inject_dir) or os.path.isdir(inject_dir):
            inject_dir = self.dir + "/" + self.object_info.mission_id().replace(" ", "") + "_ir_" + str(index) + "/"
            index = index + 1
        os.mkdir(inject_dir)
        self.__setup_logging(inject_dir)
        return inject_dir

    def __setup_logging(self, inject_dir):
        file_dir = inject_dir + "matrix.log"
        formatter = logging.Formatter('%(message)s')
        logger = logging.getLogger()
        while len(logger.handlers) > 0:
            logger.handlers.pop()
        logger.setLevel(logging.INFO)
        handler = logging.StreamHandler(sys.stdout)
        handler.setLevel(logging.INFO)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        handler = logging.FileHandler(file_dir)
        handler.setLevel(logging.INFO)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        logging.info("Setup injection directory")

    def inject(self, phases, min_period, max_period, steps_period, min_radius, max_radius, steps_radius,
               period_grid_geom="lin", radius_grid_geom="lin"):
        assert phases is not None and isinstance(phases, int) and phases > 0
        assert min_period is not None and isinstance(min_period, (int, float)) and min_period > 0
        assert max_period is not None and isinstance(max_period, (int, float)) and max_period > 0
        assert steps_period is not None and isinstance(steps_period, (int)) and steps_period > 0
        assert min_radius is not None and isinstance(min_radius, (int, float)) and min_radius > 0
        assert max_radius is not None and isinstance(max_radius, (int, float)) and max_radius > 0
        assert steps_radius is not None and isinstance(steps_radius, (int)) and steps_radius > 0
        assert max_period >= min_period
        assert max_radius >= min_radius
        inject_dir = self.retrieve_object_data()
        flux0 = self.lc_build.lc.flux.value
        time = self.lc_build.lc.time.value
        flux_err = self.lc_build.lc.flux_err.value
        period_grid = np.linspace(min_period, max_period, steps_period) if period_grid_geom == "lin" \
            else np.logspace(np.log10(min_period), np.log10(max_period), steps_period)
        radius_grid = np.linspace(min_radius, max_radius, steps_radius) if radius_grid_geom == "lin" \
            else np.logspace(np.log10(min_radius), np.log10(max_radius), steps_period)
        inject_models = []
        for period in period_grid:
            for t0 in np.arange(time[60], time[60] + period - 0.1, period / phases):
                for rplanet in radius_grid:
                    rplanet = np.around(rplanet, decimals=2) * u.R_earth
                    inject_models.append(InjectModel(inject_dir, time, flux0, flux_err, self.rstar, self.mstar, t0,
                                                     period, rplanet, self.exposure_time, self.ab))
        with Pool(processes=self.cores) as pool:
            pool.map(InjectModel.make_model, inject_models)
        return inject_dir

    def recovery(self, inject_dir, snr_threshold=5, sherlock_samples=0, detrend_ws=0,
                 transit_template='tls', run_limit=5, custom_search_algorithm=None, max_period_search=25):
        assert detrend_ws is not None and isinstance(detrend_ws, (int, float))
        assert transit_template in ('tls', 'bls')
        assert inject_dir is not None and isinstance(inject_dir, str)
        if transit_template == 'tls':
            transit_template = 'default'
        elif transit_template == 'bls':
            transit_template = 'box'
        reports_df = pd.DataFrame(columns=['period', 'radius', 'epoch', 'duration_found', 'period_found', 'epoch_found',
                                           'found', 'snr', 'sde', 'run'])
        for file in sorted(os.listdir(inject_dir)):
            file_name_matches = re.search("P([0-9]+\\.[0-9]+)+_R([0-9]+\\.[0-9]+)_([0-9]+\\.[0-9]+)\\.csv", file)
            if file_name_matches is not None:
                try:
                    period = float(file_name_matches[1])
                    r_planet = float(file_name_matches[2])
                    epoch = float(file_name_matches[3])
                    df = pd.read_csv(inject_dir + file, float_precision='round_trip', sep=',',
                                     usecols=['#time', 'flux', 'flux_err'])
                    if len(df) == 0:
                        found = True
                        snr = 20
                        sde = self.SDE_ROCHE
                        run = 1
                        duration_found = 20
                        epoch_found = 0
                        period_found = 0
                    else:
                        self.retrieve_object_data_for_recovery(inject_dir + "/", inject_dir + file)
                        found, snr, sde, run, duration_found, period_found, epoch_found = \
                            self.__search(self.lc_build.lc.time.value, self.lc_build.lc.flux.value, self.radius, self.radiusmin,
                                          self.radiusmax, self.mass, self.massmin,
                                          self.massmax, self.ab, epoch, period, self.MIN_SEARCH_PERIOD,
                                          max_period_search, snr_threshold,
                                          transit_template, detrend_ws, self.lc_build.transits_min_count,
                                          run_limit, custom_search_algorithm)
                    new_report = {"period": period, "radius": r_planet, "epoch": epoch, "found": found, "snr": snr,
                                  "sde": sde, "run": run, "duration_found": duration_found,
                                  "period_found": period_found, "epoch_found": epoch_found}
                    reports_df = reports_df.append(new_report, ignore_index=True)
                    print("P=" + str(period) + ", R=" + str(r_planet) + ", T0=" + str(epoch) + ", FOUND WAS " + str(
                        found) +
                          " WITH SNR " + str(snr) + " AND SDE " + str(sde))
                    reports_df = reports_df.sort_values(['period', 'radius', 'epoch'], ascending=[True, True, True])
                    reports_df.to_csv(inject_dir + "a_tls_report.csv", index=False)
                except Exception as e:
                    traceback.print_exc()
                    print("File not valid: " + file)
        # tls_report_df = pd.read_csv(inject_dir + "a_tls_report.csv", float_precision='round_trip', sep=',',
        #                             usecols=['period', 'radius', 'epoch', 'found', 'duration', 'snr', 'sde', 'run'])
        if sherlock_samples > 0:
            from sherlockpipe import sherlock
            from sherlockpipe.scoring.QuorumSnrBorderCorrectedSignalSelector import QuorumSnrBorderCorrectedSignalSelector

            class QuorumSnrBorderCorrectedStopWhenMatchSignalSelector(QuorumSnrBorderCorrectedSignalSelector):
                def __init__(self, strength=1, min_quorum=0, per=None, t0=None):
                    super().__init__()
                    self.strength = strength
                    self.min_quorum = min_quorum
                    self.per = per
                    self.t0 = t0

                def select(self, transit_results, snr_min, detrend_method, wl):
                    signal_selection = super(QuorumSnrBorderCorrectedStopWhenMatchSignalSelector, self) \
                        .select(transit_results, snr_min, detrend_method, wl)
                    if signal_selection.score == 0 or (
                            self.is_harmonic(signal_selection.transit_result.period, self.per) and
                            self.isRightEpoch(signal_selection.transit_result.t0, self.t0, self.per)):
                        signal_selection.score = 0
                    return signal_selection

                def is_harmonic(self, a, b, tolerance=0.05):
                    a = np.float(a)
                    b = np.float(b)
                    mod_ab = a % b
                    mod_ba = b % a
                    return (a > b and a < b * 3 + tolerance * 3 and (
                            abs(mod_ab % 1) <= tolerance or abs((b - mod_ab) % 1) <= tolerance)) or \
                           (b > a and a > b / 3 - tolerance / 3 and (
                                   abs(mod_ba % 1) <= tolerance or abs((a - mod_ba) % 1) <= tolerance))

                def isRightEpoch(self, t0, known_epoch, known_period):
                    right_epoch = False
                    for i in range(-5, 5):
                        right_epoch = right_epoch or (np.abs(t0 - known_epoch + i * known_period) < (
                                1. / 24.))
                    return right_epoch
            report = {}
            reports_df = pd.DataFrame(columns=['period', 'radius', 'epoch', 'found', 'snr', 'sde', 'run'])
            a = False
            samples_analysed = sherlock_samples
            for index, row in tls_report_df[::-1].iterrows():
                file = os.path.join(
                    'P' + str(row['period']) + '_R' + str(row['radius']) + '_' + str(row['epoch']) + '.csv')
                first_false = index > 0 and tls_report_df.iloc[index - 1]['found'] and \
                            not tls_report_df.iloc[index]['found']
                if first_false:
                    samples_analysed = 0
                elif tls_report_df.iloc[index]['found']:
                    samples_analysed = sherlock_samples
                if samples_analysed < sherlock_samples:
                    try:
                        samples_analysed = samples_analysed + 1
                        period = float(re.search("P([0-9]+\\.[0-9]+)", file)[1])
                        r_planet = float(re.search("R([0-9]+\\.[0-9]+)", file)[1])
                        epoch = float(re.search("_([0-9]+\\.[0-9]+)\\.csv", file)[1])
                        signal_selection_algorithm = QuorumSnrBorderCorrectedStopWhenMatchSignalSelector(1, 0, period,
                                                                                epoch)
                        df = pd.read_csv(inject_dir + file, float_precision='round_trip', sep=',',
                                         usecols=['#time', 'flux', 'flux_err'])
                        if len(df) == 0:
                            found = True
                            snr = 20
                            sde = 20
                            run = 1
                        else:
                            sherlock.Sherlock(False, object_infos=[MissionInputObjectInfo(self.id, inject_dir + file)]) \
                                .setup_detrend(True, True, 1.5, 4, 12, "biweight", 0.2, 1.0, 20, False,
                                               0.25, "cosine", None) \
                                .setup_transit_adjust_params(5, None, None, 10, None, None, 0.4, 14, 10,
                                                             20, 5, 5.5, 0.05, "mask", "quorum", 1, 0,
                                                             signal_selection_algorithm) \
                                .run()
                            df = pd.read_csv(self.id.replace(" ", "") + "_INP/candidates.csv", float_precision='round_trip', sep=',',
                                             usecols=['curve', 'period', 't0', 'run', 'snr', 'sde', 'rad_p',
                                                      'transits'])
                            snr = df["snr"].iloc[len(df) - 1]
                            run = df["run"].iloc[len(df) - 1]
                            sde = df["sde"].iloc[len(df) - 1]
                            per_run = 0
                            found_period = False
                            j = 0
                            for per in df["period"]:
                                if signal_selection_algorithm.is_harmonic(per, period / 2.):
                                    found_period = True
                                    t0 = df["t0"].iloc[j]
                                    break
                                j = j + 1
                            right_epoch = False
                            if found_period:
                                for i in range(-5, 5):
                                    right_epoch = right_epoch or (np.abs(t0 - epoch + i * period) < (
                                            1. / 24.))
                                    if right_epoch:
                                        snr = df["snr"].iloc[j]
                                        run = df["run"].iloc[j]
                                        sde = df["sde"].iloc[j]
                                        break
                            found = right_epoch
                        new_report = {"period": period, "radius": r_planet, "epoch": epoch, "found": found, "sde": sde,
                                      "snr": snr,
                                      "run": int(run)}
                        reports_df = reports_df.append(new_report, ignore_index=True)
                        reports_df.to_csv(inject_dir + "a_sherlock_report.csv", index=False)
                        print("P=" + str(period) + ", R=" + str(r_planet) + ", T0=" + str(epoch) + ", FOUND WAS " + str(
                            found) +
                              " WITH SNR " + str(snr) + "and SDE " + str(sde))
                    except Exception as e:
                        print(e)
                        print("File not valid: " + file)

        # If preserve parameter is not True, we remove inject files:
        if not self.preserve:
            for file in os.listdir(inject_dir):
                if file.endswith(".csv") and file.startswith("P"):
                    os.remove(inject_dir + file)

    def transit_masks(self, transit_masks, time):
        if transit_masks is None:
            transit_masks = []
        result = np.full(len(time), False)
        for mask in transit_masks:
            intransit = transit_mask(time, mask["P"], 2 * mask["D"], mask["T0"])
            result[intransit] = True

        return result

    def __clean(self, lc, detrend_period, detrend_period_method, custom_clean_algorithm):
        clean_flux = lc.flux.value
        time = lc.time.value
        if custom_clean_algorithm is not None:
            clean_flux = custom_clean_algorithm.clean(time, clean_flux)
        elif detrend_period:
            periodogram = lc.to_periodogram(minimum_period=0.05, maximum_period=15, oversample_factor=10)
            ws = self.__calculate_max_significant_period(lc, periodogram)
            clean_flux = wotan.flatten(time, clean_flux, window_length=ws, return_trend=False,
                                       method=detrend_period_method, break_tolerance=0.5)
        return clean_flux

    def __calculate_max_significant_period(self, lc, periodogram):
        #max_accepted_period = (lc.time[len(lc.time) - 1] - lc.time[0]) / 4
        max_accepted_period = np.float64(10)
        # TODO related to https://github.com/franpoz/SHERLOCK/issues/29 check whether this fits better
        max_power_index = np.argmax(periodogram.power)
        period = periodogram.period[max_power_index]
        if max_power_index > 0.0008:
            period = period.value
            logging.info("Auto-Detrend found the strong period: " + str(period) + ".")
        else:
            logging.info("Auto-Detrend did not find relevant periods.")
            period = None
        return period

    @staticmethod
    def plot_results(object_id, inject_dir, binning=1, xticks=None, yticks=None, period_grid_geom="lin",
                     radius_grid_geom="lin"):
        df = pd.read_csv(inject_dir + '/a_tls_report.csv', float_precision='round_trip', sep=',',
                         usecols=['period', 'radius', 'found', 'sde'])
        min_period = df["period"].min()
        max_period = df["period"].max()
        min_rad = df["radius"].min()
        max_rad = df["radius"].max()
        phases = len(df[df["period"] == df["period"].min()][df["radius"] == df["radius"].min()])
        phases_str = "phase" if phases == 1 else "phases"
        bin_nums = int(np.ceil(len(df["period"].unique()) / binning))
        if period_grid_geom == 'lin':
            step_period = (max_period - min_period) / (len(df["period"].unique()) - 1)
            step_period = step_period * binning
            if step_period <= 0:
                step_period = 0.1
            period_grid = np.linspace(min_period, max_period, bin_nums)\
                if max_period - min_period > 0 else np.full((1), min_period)
        else:
            period_grid = np.logspace(np.log10(min_period), np.log10(max_period), bin_nums)
        bin_nums = int(np.ceil(len(df["radius"].unique()) / binning))
        if radius_grid_geom == 'lin':
            step_radius = (max_rad - min_rad) / (len(df["radius"].unique()) - 1)
            step_radius = step_radius * binning
            if step_radius <= 0:
                step_radius = 0.1
            radius_grid = np.round(np.linspace(min_rad, max_rad, bin_nums), 2)\
                if max_rad - min_rad > 0 else np.full((1), min_rad)
        else:
            radius_grid = np.round(np.logspace(np.log10(min_rad), np.log10(max_rad), 2), bin_nums)
        f = len(period_grid) / len(radius_grid)
        bins = [period_grid, radius_grid]
        h1, x, y = np.histogram2d(df['period'][df['found'] == 1], df['radius'][df['found'] == 1], bins=bins)
        h2, x, y = np.histogram2d(df['period'][df['found'] == 0], df['radius'][df['found'] == 0], bins=bins)
        normed_hist = (100. * h1 / (h1 + h2))
        fig, ax = plt.subplots(figsize=(2.7 * 5, 5))
        im = plt.imshow(normed_hist.T, origin='lower', extent=(x[0], x[-1], y[0], y[-1]), interpolation='none',
                        aspect='auto', cmap='viridis', vmin=0, vmax=100, rasterized=True)
        plt.colorbar(im, label='Recovery rate (%)')
        plt.xlabel('Injected period (days)')
        plt.ylabel(r'Injected radius (R$_\oplus$)')
        ax.set_title(object_id + " - P/R recovery (" + str(phases) + " " + phases_str + ")")
        if xticks is not None:
            plt.xticks(xticks)
        else:
            period_ticks_decimals = MATRIX.num_of_zeros(max_period - min_period) + 1
            plot_bins = 10 if 10 < len(period_grid) else len(period_grid)
            plt.locator_params(axis="x", nbins=plot_bins)
            ax.xaxis.set_major_formatter(FormatStrFormatter('%.' + str(period_ticks_decimals) + 'f'))
        if yticks is not None:
            plt.xticks(yticks)
        plt.savefig(inject_dir + '/inj-rec.png', bbox_inches='tight', dpi=200)
        plt.close()

    @staticmethod
    def plot_diff(object_id, inject_dir1, inject_dir2, output_dir, binning=1, xticks=None, yticks=None,
                  period_grid_geom="lin", radius_grid_geom="lin"):
        df1 = pd.read_csv(inject_dir1 + '/a_tls_report.csv', float_precision='round_trip', sep=',',
                         usecols=['period', 'radius', 'found', 'sde'])
        df2 = pd.read_csv(inject_dir2 + '/a_tls_report.csv', float_precision='round_trip', sep=',',
                         usecols=['period', 'radius', 'found', 'sde'])
        min_period = df1["period"].min()
        max_period = df1["period"].max()
        min_rad = df1["radius"].min()
        max_rad = df1["radius"].max()
        phases = len(df1[df1["period"] == df1["period"].min()][df1["radius"] == df1["radius"].min()])
        phases_str = "phase" if phases == 1 else "phases"
        bin_nums = int(np.ceil(len(df1["period"].unique()) / binning))
        if period_grid_geom == 'lin':
            step_period = (max_period - min_period) / (len(df1["period"].unique()) - 1)
            step_period = step_period * binning
            if step_period <= 0:
                step_period = 0.1
            period_grid = np.linspace(min_period, max_period, bin_nums)\
                if max_period - min_period > 0 else np.full((1), min_period)
        else:
            period_grid = np.logspace(np.log10(min_period), np.log10(max_period), bin_nums)
        bin_nums = int(np.ceil(len(df1["radius"].unique()) / binning))
        if radius_grid_geom == 'lin':
            step_radius = (max_rad - min_rad) / (len(df1["radius"].unique()) - 1)
            step_radius = step_radius * binning
            if step_radius <= 0:
                step_radius = 0.1
            radius_grid = np.round(np.linspace(min_rad, max_rad, bin_nums), 2)\
                if max_rad - min_rad > 0 else np.full((1), min_rad)
        else:
            radius_grid = np.round(np.logspace(np.log10(min_rad), np.log10(max_rad), 2), bin_nums)
        f = len(period_grid) / len(radius_grid)
        bins = [period_grid, radius_grid]
        h11, x1, y1 = np.histogram2d(df1['period'][df1['found'] == 1], df1['radius'][df1['found'] == 1], bins=bins)
        h12, x1, y1 = np.histogram2d(df1['period'][df1['found'] == 0], df1['radius'][df1['found'] == 0], bins=bins)
        h21, x2, y2 = np.histogram2d(df2['period'][df2['found'] == 1], df2['radius'][df2['found'] == 1], bins=bins)
        h22, x2, y2 = np.histogram2d(df2['period'][df2['found'] == 0], df2['radius'][df2['found'] == 0], bins=bins)
        h1 = h11 - h12
        h2 = h21 - h22
        normed_hist1 = phases * h11 / (h11 + h12)
        normed_hist2 = phases * h21 / (h21 + h22)
        normed_hist = normed_hist1 - normed_hist2
        fig, ax = plt.subplots(figsize=(2.7 * 5, 5))
        im = plt.imshow(normed_hist.T, origin='lower', extent=(x1[0], x1[-1], y1[0], y1[-1]), interpolation='none',
                        aspect='auto', cmap='viridis', vmin=-phases, vmax=phases, rasterized=True)
        plt.colorbar(im, label='# Found samples diff.')
        plt.xlabel('Injected period (days)')
        plt.ylabel(r'Injected radius (R$_\oplus$)')
        ax.set_title(object_id + " - P/R recovery diff(" + str(phases) + " " + phases_str + ")")
        if xticks is not None:
            plt.xticks(xticks)
        else:
            period_ticks_decimals = MATRIX.num_of_zeros(max_period - min_period) + 1
            plot_bins = 10 if 10 < len(period_grid) else len(period_grid)
            plt.locator_params(axis="x", nbins=plot_bins)
            ax.xaxis.set_major_formatter(FormatStrFormatter('%.' + str(period_ticks_decimals) + 'f'))
        if yticks is not None:
            plt.xticks(yticks)
        plt.savefig(output_dir + '/inj-rec-diff.png', bbox_inches='tight', dpi=200)
        plt.close()

    def __search(self, time, flux, rstar, rstar_min, rstar_max, mass, mstar_min, mstar_max, ab, epoch,
                 period, min_period, max_period, min_snr, transit_template, ws, transits_min_count,
                 run_limit, custom_search_algorithm):
        tls_period_grid = self.__calculate_period_grid(time, min_period, max_period, 3, self.star_info,
                                                   transits_min_count)
        if custom_search_algorithm is not None:
            return custom_search_algorithm.search(time, flux, rstar, rstar_min, rstar_max, mass, mstar_min, mstar_max,
                                                ab, epoch, period, min_period, max_period, min_snr, self.cores,
                                                transit_template, ws, transits_min_count, run_limit)
        else:
            return self.__tls_search(time, flux, rstar, rstar_min, rstar_max, mass, mstar_min, mstar_max, ab, epoch,
                     period, min_period, max_period, min_snr, self.cores, transit_template, ws, transits_min_count,
                     run_limit, tls_period_grid)

    def __tls_search(self, time, flux, rstar, rstar_min, rstar_max, mass, mstar_min, mstar_max, ab, epoch,
                     period, min_period, max_period, min_snr, cores, transit_template, ws, transits_min_count,
                     run_limit, tls_period_grid):
        snr = 1e12
        found_signal = False
        time, flux = cleaned_array(time, flux)
        run = 0
        if ws > 0:
            flux = wotan.flatten(time, flux, window_length=ws, return_trend=False, method='biweight', break_tolerance=0.5)
        while snr >= min_snr and not found_signal and (run_limit > 0 and run < run_limit):
            model = transitleastsquares(time, flux)
            # R_starx = rstar / u.R_sun
            results = model.power(u=ab,
                                  R_star=rstar,  # rstar/u.R_sun,
                                  R_star_min=rstar_min,  # rstar_min/u.R_sun,
                                  R_star_max=rstar_max,  # rstar_max/u.R_sun,
                                  M_star=mass,  # mstar/u.M_sun,
                                  M_star_min=mstar_min,  # mstar_min/u.M_sun,
                                  M_star_max=mstar_max,  # mstar_max/u.M_sun,
                                  period_min=min_period,
                                  period_max=max_period,
                                  n_transits_min=transits_min_count,
                                  show_progress_bar=False,
                                  use_threads=cores,
                                  transit_template=transit_template,
                                  period_grid=tls_period_grid
                                  )
            snr = results.snr
            if results.snr >= min_snr:
                intransit_result = transit_mask(time, results.period, 2 * results.duration, results.T0)
                time = time[~intransit_result]
                flux = flux[~intransit_result]
                time, flux = cleaned_array(time, flux)
                right_period = self.__is_multiple_of(results.period, period / 2.)
                right_epoch = False
                for tt in results.transit_times:
                    for i in range(-5, 5):
                        right_epoch = right_epoch or (np.abs(tt - epoch + i * period) < (1. / 24.))
                #            right_depth   = (np.abs(np.sqrt(1.-results.depth)*rstar - rplanet)/rplanet < 0.05) #check if the depth matches
                if right_period and right_epoch:
                    found_signal = True
                    break
            run = run + 1
        return found_signal, results.snr, results.SDE, run, results.duration, results.period, results.T0

    def __equal(self, a, b, tolerance=0.01):
        return np.abs(a - b) < tolerance

    def __calculate_period_grid(self, time, min_period, max_period, oversampling, star_info, transits_min_count):
        dif = time[1:] - time[:-1]
        jumps = np.where(dif > 1)[0]
        jumps = np.append(jumps, len(time))
        previous_jump_index = 0
        time_span_all_sectors = 0
        for jumpIndex in jumps:
            time_chunk = time[previous_jump_index + 1:jumpIndex]  # ignoring first measurement as could be the last from the previous chunk
            time_span_all_sectors = time_span_all_sectors + (time_chunk[-1] - time_chunk[0])
            previous_jump_index = jumpIndex
        return DefaultTransitTemplateGenerator() \
            .period_grid(star_info.radius, star_info.mass, time[-1] - time[0], min_period, max_period, oversampling,
                         transits_min_count, time_span_all_sectors)

    @staticmethod
    def num_of_zeros(n):
        if n.is_integer():
            return 0
        s = '{:.16f}'.format(n).split('.')[1]
        return len(s) - len(s.lstrip('0'))

    def __is_multiple_of(self, a, b, tolerance=0.05):
        a = np.float(a)
        b = np.float(b)
        mod_ab = a % b
        mod_ba = b % a
        return (a > b and a < b * 3 + tolerance * 3 and (
                    abs(mod_ab % 1) <= tolerance or abs((b - mod_ab) % 1) <= tolerance)) or \
               (b > a and a > b / 3 - tolerance / 3 and (
                           abs(mod_ba % 1) <= tolerance or abs((a - mod_ba) % 1) <= tolerance))


import torch
import numpy as np
import os
from skimage import draw
from geotnf.transformation import GeometricTnf
from geotnf.flow import th_sampling_grid_to_np_flow, write_flo_file
import torch.nn.functional as F
from data.pf_dataset import PFDataset, PFPascalDataset
from data.caltech_dataset import CaltechDataset
from torch.autograd import Variable
from geotnf.point_tnf import PointTnf, PointsToUnitCoords, PointsToPixelCoords
from util.py_util import create_file_path
from model.loss import WeakInlierCount, TwoStageWeakInlierCount

def compute_metric(metric,model,dataset,dataloader,batch_tnf,batch_size,two_stage=True,do_aff=False,do_tps=False,args=None):
    # Initialize stats
    N=len(dataset)
    stats={}
    # decide which results should be computed aff/tps/aff+tps
    if two_stage or do_aff:
        stats['aff']={}
    if not two_stage and do_tps:
        stats['tps']={}
    if two_stage:
        stats['aff_tps']={}
    # choose metric function and metrics to compute
    if metric=='pck':
        metrics = ['pck']
        metric_fun = pck_metric
    if metric=='dist':
        metrics = ['dist']
        metric_fun = point_dist_metric
    elif metric=='area':
        metrics = ['intersection_over_union',
                   'label_transfer_accuracy',
                   'localization_error']
        metric_fun = area_metrics
    elif metric=='pascal_parts':
        metrics = ['intersection_over_union','pck']
        metric_fun = pascal_parts_metrics
    elif metric=='flow':
        metrics = ['flow']
        metric_fun = flow_metrics
    elif metric=='inlier_count':
        metrics = ['inlier_count']
        metric_fun = inlier_count
        model.return_correlation = True
    # initialize vector for storing results for each metric
    for key in stats.keys():
        for metric in metrics:
            stats[key][metric] = np.zeros((N,1))

    # Compute
    for i, batch in enumerate(dataloader):
        batch = batch_tnf(batch)
        batch_start_idx=batch_size*i
        batch_end_idx=np.minimum(batch_start_idx+batch_size,N)

        model.eval()
        theta_aff=None
        theta_tps=None
        theta_aff_tps=None

        if two_stage:
            if model.return_correlation==False:
                theta_aff,theta_aff_tps=model(batch)
            else:
                theta_aff,theta_aff_tps,corr_aff,corr_aff_tps=model(batch)
        elif do_aff:
            theta_aff=model(batch)
            if isinstance(theta_aff,tuple):
                theta_aff=theta_aff[0]
        elif do_tps:
            theta_tps=model(batch)
            if isinstance(theta_tps,tuple):
                theta_tps=theta_tps[0]

        if metric=='inlier_count':
            stats = inlier_count(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,corr_aff,corr_aff_tps,stats,args)
        elif metric_fun is not None:
            stats = metric_fun(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,stats,args)

        print('Batch: [{}/{} ({:.0f}%)]'.format(i, len(dataloader), 100. * i / len(dataloader)))

    # Print results
    if metric == 'flow':
        print('Flow files have been saved to '+args.flow_output_dir)
        return stats

    for key in stats.keys():
        print('=== Results '+key+' ===')
        for metric in metrics:
            # print per-class brakedown for PFPascal, or caltech
            if isinstance(dataset,PFPascalDataset):
                N_cat = int(np.max(dataset.category))
                for c in range(N_cat):
                    cat_idx = np.nonzero(dataset.category==c+1)[0]
                    print(dataset.category_names[c].ljust(15)+': ','{:.2%}'.format(np.mean(stats[key][metric][cat_idx])))

            # print mean value
            results=stats[key][metric]
            good_idx = np.flatnonzero((results!=-1) * ~np.isnan(results))
            print('Total: '+str(results.size))
            print('Valid: '+str(good_idx.size))
            filtered_results = results[good_idx]
            print(metric+':','{:.2%}'.format(np.mean(filtered_results)))

        print('\n')

    return stats

def pck(source_points,warped_points,L_pck,alpha=0.1):
    # compute precentage of correct keypoints
    batch_size=source_points.size(0)
    pck=torch.zeros((batch_size))
    for i in range(batch_size):
        p_src = source_points[i,:]
        p_wrp = warped_points[i,:]
        N_pts = torch.sum(torch.ne(p_src[0,:],-1)*torch.ne(p_src[1,:],-1))
        point_distance = torch.pow(torch.sum(torch.pow(p_src[:,:N_pts]-p_wrp[:,:N_pts],2),0),0.5)
        L_pck_mat = L_pck[i].expand_as(point_distance)
        correct_points = torch.le(point_distance,L_pck_mat*alpha)
        pck[i]=torch.mean(correct_points.float())
    return pck

def mean_dist(source_points,warped_points,L_pck):
    # compute precentage of correct keypoints
    batch_size=source_points.size(0)
    dist=torch.zeros((batch_size))
    for i in range(batch_size):
        p_src = source_points[i,:]
        p_wrp = warped_points[i,:]
        N_pts = torch.sum(torch.ne(p_src[0,:],-1)*torch.ne(p_src[1,:],-1))
        point_distance = torch.pow(torch.sum(torch.pow(p_src[:,:N_pts]-p_wrp[:,:N_pts],2),0),0.5)
        L_pck_mat = L_pck[i].expand_as(point_distance)
        dist[i]=torch.mean(torch.div(point_distance,L_pck_mat))
    return dist

def point_dist_metric(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,stats,args,use_cuda=True):
    do_aff = theta_aff is not None
    do_tps = theta_tps is not None
    do_aff_tps = theta_aff_tps is not None

    source_im_size = batch['source_im_size']
    target_im_size = batch['target_im_size']

    source_points = batch['source_points']
    target_points = batch['target_points']

    # Instantiate point transformer
    pt = PointTnf(use_cuda=use_cuda,
                  tps_reg_factor=args.tps_reg_factor)

    # warp points with estimated transformations
    target_points_norm = PointsToUnitCoords(target_points,target_im_size)

    if do_aff:
        # do affine only
        warped_points_aff_norm = pt.affPointTnf(theta_aff,target_points_norm)
        warped_points_aff = PointsToPixelCoords(warped_points_aff_norm,source_im_size)

    if do_tps:
        # do tps only
        warped_points_tps_norm = pt.tpsPointTnf(theta_tps,target_points_norm)
        warped_points_tps = PointsToPixelCoords(warped_points_tps_norm,source_im_size)

    if do_aff_tps:
        # do tps+affine
        warped_points_aff_tps_norm = pt.tpsPointTnf(theta_aff_tps,target_points_norm)
        warped_points_aff_tps_norm = pt.affPointTnf(theta_aff,warped_points_aff_tps_norm)
        warped_points_aff_tps = PointsToPixelCoords(warped_points_aff_tps_norm,source_im_size)

    L_pck = batch['L_pck'].data

    current_batch_size=batch['source_im_size'].size(0)
    indices = range(batch_start_idx,batch_start_idx+current_batch_size)

#    import pdb; pdb.set_trace()

    if do_aff:
        dist_aff = mean_dist(source_points.data, warped_points_aff.data, L_pck)

    if do_tps:
        dist_tps = mean_dist(source_points.data, warped_points_tps.data, L_pck)

    if do_aff_tps:
        dist_aff_tps = mean_dist(source_points.data, warped_points_aff_tps.data, L_pck)

    if do_aff:
        stats['aff']['dist'][indices] = dist_aff.unsqueeze(1).cpu().numpy()
    if do_tps:
        stats['tps']['dist'][indices] = dist_tps.unsqueeze(1).cpu().numpy()
    if do_aff_tps:
        stats['aff_tps']['dist'][indices] = dist_aff_tps.unsqueeze(1).cpu().numpy()

    return stats

def inlier_count(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,corr_aff,corr_aff_tps,stats,args,use_cuda=True):
    inliersComposed = TwoStageWeakInlierCount(use_cuda=torch.cuda.is_available(),dilation_filter=0,normalize_inlier_count=True)
    inliers_comp = inliersComposed(matches=corr_aff,theta_aff=theta_aff,theta_aff_tps=theta_aff_tps)
    current_batch_size=batch['source_im_size'].size(0)
    indices = range(batch_start_idx,batch_start_idx+current_batch_size)

    stats['aff_tps']['inlier_count'][indices] = inliers_comp.unsqueeze(1).cpu().data.numpy()

    return stats

def pck_metric(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,stats,args,use_cuda=True):
    alpha = args.pck_alpha
    do_aff = theta_aff is not None
    do_tps = theta_tps is not None
    do_aff_tps = theta_aff_tps is not None

    source_im_size = batch['source_im_size']
    target_im_size = batch['target_im_size']

    source_points = batch['source_points']
    target_points = batch['target_points']

    # Instantiate point transformer
    pt = PointTnf(use_cuda=use_cuda,
                  tps_reg_factor=args.tps_reg_factor)

    # warp points with estimated transformations
    target_points_norm = PointsToUnitCoords(target_points,target_im_size)

    if do_aff:
        # do affine only
        warped_points_aff_norm = pt.affPointTnf(theta_aff,target_points_norm)
        warped_points_aff = PointsToPixelCoords(warped_points_aff_norm,source_im_size)

    if do_tps:
        # do tps only
        warped_points_tps_norm = pt.tpsPointTnf(theta_tps,target_points_norm)
        warped_points_tps = PointsToPixelCoords(warped_points_tps_norm,source_im_size)

    if do_aff_tps:
        # do tps+affine
        warped_points_aff_tps_norm = pt.tpsPointTnf(theta_aff_tps,target_points_norm)
        warped_points_aff_tps_norm = pt.affPointTnf(theta_aff,warped_points_aff_tps_norm)
        warped_points_aff_tps = PointsToPixelCoords(warped_points_aff_tps_norm,source_im_size)

    L_pck = batch['L_pck'].data

    current_batch_size=batch['source_im_size'].size(0)
    indices = range(batch_start_idx,batch_start_idx+current_batch_size)

    # import pdb; pdb.set_trace()

    if do_aff:
        pck_aff = pck(source_points.data, warped_points_aff.data, L_pck, alpha)

    if do_tps:
        pck_tps = pck(source_points.data, warped_points_tps.data, L_pck, alpha)

    if do_aff_tps:
        pck_aff_tps = pck(source_points.data, warped_points_aff_tps.data, L_pck, alpha)

    if do_aff:
        stats['aff']['pck'][indices] = pck_aff.unsqueeze(1).cpu().numpy()
    if do_tps:
        stats['tps']['pck'][indices] = pck_tps.unsqueeze(1).cpu().numpy()
    if do_aff_tps:
        stats['aff_tps']['pck'][indices] = pck_aff_tps.unsqueeze(1).cpu().numpy()

    return stats

def pascal_parts_metrics(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,stats,args,use_cuda=True):
    do_aff = theta_aff is not None
    do_tps = theta_tps is not None
    do_aff_tps = theta_aff_tps is not None

    batch_size=batch['source_im_size'].size(0)
    for b in range(batch_size):
        idx = batch_start_idx+b
        h_src = int(batch['source_im_size'][b,0].data.cpu().numpy())
        w_src = int(batch['source_im_size'][b,1].data.cpu().numpy())
        h_tgt = int(batch['target_im_size'][b,0].data.cpu().numpy())
        w_tgt = int(batch['target_im_size'][b,1].data.cpu().numpy())

        # do pck
        if batch['keypoint_A'][b].size!=0:
            src_points = Variable(torch.FloatTensor(batch['keypoint_A'][b])).unsqueeze(0)
            tgt_points = Variable(torch.FloatTensor(batch['keypoint_B'][b])).unsqueeze(0)
            L_pck = Variable(torch.FloatTensor([batch['L_pck'][b]])).unsqueeze(1)
            if use_cuda:
                src_points=src_points.cuda()
                tgt_points=tgt_points.cuda()
                L_pck = L_pck.cuda()

            batch_b = {'source_im_size': batch['source_im_size'][b,:].unsqueeze(0),
                       'target_im_size': batch['target_im_size'][b,:].unsqueeze(0),
                       'source_points':  src_points,
                       'target_points': tgt_points,
                       'L_pck': L_pck}
            args.pck_alpha = 0.05
            stats = pck_metric(batch_b,
                               idx,
                               theta_aff[b,:].unsqueeze(0) if do_aff else None,
                               theta_tps[b,:].unsqueeze(0) if do_tps else None,
                               theta_aff_tps[b,:].unsqueeze(0) if do_aff_tps else None,
                               stats,args,use_cuda)
        else:
            if do_aff:
                stats['aff']['pck'][idx] = -1
            if do_tps:
                stats['tps']['pck'][idx] = -1
            if do_aff_tps:
                stats['aff_tps']['pck'][idx] = -1

        # do area
        source_mask = Variable(torch.FloatTensor(batch['part_A'][b].astype(np.float32)).unsqueeze(0).transpose(2,3).transpose(1,2))
        target_mask = Variable(torch.FloatTensor(batch['part_B'][b].astype(np.float32)).unsqueeze(0).transpose(2,3).transpose(1,2))

        if use_cuda:
            source_mask = source_mask.cuda()
            target_mask = target_mask.cuda()

        grid_aff,grid_tps,grid_aff_tps=theta_to_sampling_grid(h_tgt,w_tgt,
                                                              theta_aff[b,:] if do_aff else None,
                                                              theta_tps[b,:] if do_tps else None,
                                                              theta_aff_tps[b,:] if do_aff_tps else None,
                                                              use_cuda=use_cuda,
                                                              tps_reg_factor=args.tps_reg_factor)

        if do_aff:
            warped_mask_aff = F.grid_sample(source_mask, grid_aff)
            stats['aff']['intersection_over_union'][idx] = intersection_over_union(warped_mask_aff,target_mask)
        if do_tps:
            warped_mask_tps = F.grid_sample(source_mask, grid_tps)
            stats['tps']['intersection_over_union'][idx] = intersection_over_union(warped_mask_tps,target_mask)
        if do_aff_tps:
            warped_mask_aff_tps = F.grid_sample(source_mask, grid_aff_tps)
            stats['aff_tps']['intersection_over_union'][idx] = intersection_over_union(warped_mask_aff_tps,target_mask)

    return stats

def area_metrics(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,stats,args,use_cuda=True):
    do_aff = theta_aff is not None
    do_tps = theta_tps is not None
    do_aff_tps = theta_aff_tps is not None

    batch_size=batch['source_im_size'].size(0)

    pt=PointTnf(use_cuda=use_cuda)

    for b in range(batch_size):
        h_src = int(batch['source_im_size'][b,0].data.cpu().numpy())
        w_src = int(batch['source_im_size'][b,1].data.cpu().numpy())
        h_tgt = int(batch['target_im_size'][b,0].data.cpu().numpy())
        w_tgt = int(batch['target_im_size'][b,1].data.cpu().numpy())

        target_mask_np,target_mask = poly_str_to_mask(batch['target_polygon'][0][b],
                                                      batch['target_polygon'][1][b],
                                                      h_tgt,w_tgt,use_cuda=use_cuda)

        source_mask_np,source_mask = poly_str_to_mask(batch['source_polygon'][0][b],
                                                      batch['source_polygon'][1][b],
                                                      h_src,w_src,use_cuda=use_cuda)

        grid_X,grid_Y = np.meshgrid(np.linspace(-1,1,w_tgt),np.linspace(-1,1,h_tgt))
        grid_X = torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3)
        grid_Y = torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3)
        grid_X = Variable(grid_X,requires_grad=False)
        grid_Y = Variable(grid_Y,requires_grad=False)

        if use_cuda:
            grid_X = grid_X.cuda()
            grid_Y = grid_Y.cuda()

        grid_X_vec = grid_X.view(1,1,-1)
        grid_Y_vec = grid_Y.view(1,1,-1)

        grid_XY_vec = torch.cat((grid_X_vec,grid_Y_vec),1)

        def pointsToGrid (x,h_tgt=h_tgt,w_tgt=w_tgt): return x.contiguous().view(1,2,h_tgt,w_tgt).transpose(1,2).transpose(2,3)

        idx = batch_start_idx+b

        if do_aff:
            grid_aff = pointsToGrid(pt.affPointTnf(theta_aff[b,:].unsqueeze(0),grid_XY_vec))
            warped_mask_aff = F.grid_sample(source_mask, grid_aff)
            flow_aff = th_sampling_grid_to_np_flow(source_grid=grid_aff,h_src=h_src,w_src=w_src)

            stats['aff']['intersection_over_union'][idx] = intersection_over_union(warped_mask_aff,target_mask)
            stats['aff']['label_transfer_accuracy'][idx] = label_transfer_accuracy(warped_mask_aff,target_mask)
            stats['aff']['localization_error'][idx] = localization_error(source_mask_np, target_mask_np, flow_aff)
        if do_tps:
            grid_tps = pointsToGrid(pt.tpsPointTnf(theta_tps[b,:].unsqueeze(0),grid_XY_vec))
            warped_mask_tps = F.grid_sample(source_mask, grid_tps)
            flow_tps = th_sampling_grid_to_np_flow(source_grid=grid_tps,h_src=h_src,w_src=w_src)

            stats['tps']['intersection_over_union'][idx] = intersection_over_union(warped_mask_tps,target_mask)
            stats['tps']['label_transfer_accuracy'][idx] = label_transfer_accuracy(warped_mask_tps,target_mask)
            stats['tps']['localization_error'][idx] = localization_error(source_mask_np, target_mask_np, flow_tps)
        if do_aff_tps:
            grid_aff_tps = pointsToGrid(pt.affPointTnf(theta_aff[b,:].unsqueeze(0),pt.tpsPointTnf(theta_aff_tps[b,:].unsqueeze(0),grid_XY_vec)))
            warped_mask_aff_tps = F.grid_sample(source_mask, grid_aff_tps)
            flow_aff_tps = th_sampling_grid_to_np_flow(source_grid=grid_aff_tps,h_src=h_src,w_src=w_src)

            stats['aff_tps']['intersection_over_union'][idx] = intersection_over_union(warped_mask_aff_tps,target_mask)
            stats['aff_tps']['label_transfer_accuracy'][idx] = label_transfer_accuracy(warped_mask_aff_tps,target_mask)
            stats['aff_tps']['localization_error'][idx] = localization_error(source_mask_np, target_mask_np, flow_aff_tps)

    return stats

def flow_metrics(batch,batch_start_idx,theta_aff,theta_tps,theta_aff_tps,stats,args,use_cuda=True):
    result_path=args.flow_output_dir

    do_aff = theta_aff is not None
    do_tps = theta_tps is not None
    do_aff_tps = theta_aff_tps is not None

    pt=PointTnf(use_cuda=use_cuda)

    batch_size=batch['source_im_size'].size(0)
    for b in range(batch_size):
        h_src = int(batch['source_im_size'][b,0].data.cpu().numpy())
        w_src = int(batch['source_im_size'][b,1].data.cpu().numpy())
        h_tgt = int(batch['target_im_size'][b,0].data.cpu().numpy())
        w_tgt = int(batch['target_im_size'][b,1].data.cpu().numpy())

        grid_X,grid_Y = np.meshgrid(np.linspace(-1,1,w_tgt),np.linspace(-1,1,h_tgt))
        grid_X = torch.FloatTensor(grid_X).unsqueeze(0).unsqueeze(3)
        grid_Y = torch.FloatTensor(grid_Y).unsqueeze(0).unsqueeze(3)
        grid_X = Variable(grid_X,requires_grad=False)
        grid_Y = Variable(grid_Y,requires_grad=False)
        if use_cuda:
            grid_X = grid_X.cuda()
            grid_Y = grid_Y.cuda()

        grid_X_vec = grid_X.view(1,1,-1)
        grid_Y_vec = grid_Y.view(1,1,-1)

        grid_XY_vec = torch.cat((grid_X_vec,grid_Y_vec),1)

        def pointsToGrid (x,h_tgt=h_tgt,w_tgt=w_tgt): return x.contiguous().view(1,2,h_tgt,w_tgt).transpose(1,2).transpose(2,3)

        idx = batch_start_idx+b

        if do_aff:
            grid_aff = pointsToGrid(pt.affPointTnf(theta_aff[b,:].unsqueeze(0),grid_XY_vec))
            flow_aff = th_sampling_grid_to_np_flow(source_grid=grid_aff,h_src=h_src,w_src=w_src)
            flow_aff_path = os.path.join(result_path,'aff',batch['flow_path'][b])
            create_file_path(flow_aff_path)
            write_flo_file(flow_aff,flow_aff_path)
        if do_tps:
            grid_tps = pointsToGrid(pt.tpsPointTnf(theta_tps[b,:].unsqueeze(0),grid_XY_vec))
            flow_tps = th_sampling_grid_to_np_flow(source_grid=grid_tps,h_src=h_src,w_src=w_src)
            flow_tps_path = os.path.join(result_path,'tps',batch['flow_path'][b])
            create_file_path(flow_tps_path)
            write_flo_file(flow_tps,flow_tps_path)
        if do_aff_tps:
            grid_aff_tps = pointsToGrid(pt.affPointTnf(theta_aff[b,:].unsqueeze(0),pt.tpsPointTnf(theta_aff_tps[b,:].unsqueeze(0),grid_XY_vec)))
            flow_aff_tps = th_sampling_grid_to_np_flow(source_grid=grid_aff_tps,h_src=h_src,w_src=w_src)
            flow_aff_tps_path = os.path.join(result_path,'aff_tps',batch['flow_path'][b])
            create_file_path(flow_aff_tps_path)
            write_flo_file(flow_aff_tps,flow_aff_tps_path)

        idx = batch_start_idx+b
    return stats

def poly_to_mask(vertex_row_coords, vertex_col_coords, shape):
    fill_row_coords, fill_col_coords = draw.polygon(vertex_row_coords, vertex_col_coords, shape)
    mask = np.zeros(shape, dtype=np.bool)
    mask[fill_row_coords, fill_col_coords] = True
    return mask

def poly_str_to_mask(poly_x_str,poly_y_str,out_h,out_w,use_cuda=True):
    polygon_x = np.fromstring(poly_x_str,sep=',')
    polygon_y = np.fromstring(poly_y_str,sep=',')
    mask_np = poly_to_mask(vertex_col_coords=polygon_x,
                               vertex_row_coords=polygon_y,shape=[out_h,out_w])
    mask = Variable(torch.FloatTensor(mask_np.astype(np.float32)).unsqueeze(0).unsqueeze(0))
    if use_cuda:
        mask = mask.cuda()
    return (mask_np,mask)

#def intersection_over_union(warped_mask,target_mask):
#    return torch.sum(warped_mask.data.gt(0.5) & target_mask.data.gt(0.5))/torch.sum(warped_mask.data.gt(0.5) | target_mask.data.gt(0.5))
def intersection_over_union(warped_mask,target_mask):
    relative_part_weight = torch.sum(torch.sum(target_mask.data.gt(0.5).float(),2,True),3,True)/torch.sum(target_mask.data.gt(0.5).float())
    part_iou = torch.sum(torch.sum((warped_mask.data.gt(0.5) & target_mask.data.gt(0.5)).float(),2,True),3,True)/torch.sum(torch.sum((warped_mask.data.gt(0.5) | target_mask.data.gt(0.5)).float(),2,True),3,True)
    weighted_iou = torch.sum(torch.mul(relative_part_weight,part_iou))
    return weighted_iou

def label_transfer_accuracy(warped_mask,target_mask):
    return torch.mean((warped_mask.data.gt(0.5) == target_mask.data.gt(0.5)).double())

def localization_error(source_mask_np, target_mask_np, flow_np):
    h_tgt, w_tgt = target_mask_np.shape[0],target_mask_np.shape[1]
    h_src, w_src = source_mask_np.shape[0],source_mask_np.shape[1]

    # initial pixel positions x1,y1 in target image
    x1, y1 = np.meshgrid(range(1,w_tgt+1), range(1,h_tgt+1))
    # sampling pixel positions x2,y2
    x2 = x1 + flow_np[:,:,0]
    y2 = y1 + flow_np[:,:,1]

    # compute in-bound coords for each image
    in_bound = (x2 >= 1) & (x2 <= w_src) & (y2 >= 1) & (y2 <= h_src)
    row,col = np.where(in_bound)
    row_1=y1[row,col].flatten().astype(np.int)-1
    col_1=x1[row,col].flatten().astype(np.int)-1
    row_2=y2[row,col].flatten().astype(np.int)-1
    col_2=x2[row,col].flatten().astype(np.int)-1

    # compute relative positions
    target_loc_x,target_loc_y = obj_ptr(target_mask_np)
    source_loc_x,source_loc_y = obj_ptr(source_mask_np)
    x1_rel=target_loc_x[row_1,col_1]
    y1_rel=target_loc_y[row_1,col_1]
    x2_rel=source_loc_x[row_2,col_2]
    y2_rel=source_loc_y[row_2,col_2]

    # compute localization error
    loc_err = np.mean(np.abs(x1_rel-x2_rel)+np.abs(y1_rel-y2_rel))

    return loc_err

def obj_ptr(mask):
    # computes images of normalized coordinates around bounding box
    # kept function name from DSP code
    h,w = mask.shape[0],mask.shape[1]
    y, x = np.where(mask>0.5)
    left = np.min(x);
    right = np.max(x);
    top = np.min(y);
    bottom = np.max(y);
    fg_width = right-left + 1;
    fg_height = bottom-top + 1;
    x_image,y_image = np.meshgrid(range(1,w+1), range(1,h+1));
    x_image = (x_image - left)/fg_width;
    y_image = (y_image - top)/fg_height;
    return (x_image,y_image)


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

"""Python dynamics helpers.

Provides a simplified control-oriented model of the wing and tether.
"""
import collections
import copy

from makani.analysis.control import actuator_util
from makani.analysis.control import catenary
from makani.analysis.control import geometry
from makani.analysis.control import type_util
from makani.control import control_types
from makani.control import system_types
from makani.sim.physics import physics
import numpy as np

# Structure storing a force, moment, and position at which that force
# is applied.
ForceMomentPos = collections.namedtuple('ForceMomentPos',
                                        ['force', 'moment', 'pos'])

# Structure for storing forces and moment.
ForceMoment = collections.namedtuple('ForceMoment', ['force', 'moment'])

# Structure representing the inputs to the wing.
#   thrust: Motor thrust [N] (1-by-1 np.matrix).
#   motor_moment: Motor moments [N-m] (3-by-1 np.matrix).
#   flaps: Flaps [rad] (kNumFlaps-by-1 np.matrix).
#   wind_g: Wind speed [m/s] in ground coordinates (3-by-1 np.matrix).
WingInputs = type_util.MakeNamedVectorClass(  # pylint: disable=invalid-name
    'WingInputs', [('thrust', range(0, 1)),
                   ('motor_moment', range(1, 4)),
                   ('flaps', range(4, 4 + system_types.kNumFlaps)),
                   ('wind_g', range(4 + system_types.kNumFlaps,
                                    7 + system_types.kNumFlaps))])

class WingState(type_util.MakeStateClass(
    'WingState', [('omega_b', range(0, 3)),
                  ('dcm_g2b', range(3, 6)),
                  ('wing_vel_g', range(6, 9)),
                  ('wing_pos_g', range(9, 12))])):
  """Class representing the state of the wing.

  Attributes:
    omega_b: Body angular rates.
    dcm_g2b: Ground to body rotation DCM.  Increments in the DCM are represented
        by an Euler vector.
    wing_vel_g: Velocity of the wing in ground coordinates.
    wing_pos_g: Position of the wing in ground coordinates.
  """

  def Increment(self, tangent, step=1.0):
    """Return a state evolved from this state along a tangent direction.

    Args:
      tangent: A WingState.Tangent along which to move.
      step: A scaling of how far to move.

    Returns:
      A new WingState.
    """
    return WingState(omega_b=self.omega_b + step * tangent.domega_b,
                     dcm_g2b=(geometry.AxisToDcm(step * tangent.ddcm_g2b)
                              * self.dcm_g2b),
                     wing_vel_g=self.wing_vel_g + step * tangent.dwing_vel_g,
                     wing_pos_g=self.wing_pos_g + step * tangent.dwing_pos_g)

  def Difference(self, other_state):
    """Inverse operation of Increment with a step size of 1.0."""
    return WingState.Tangent(
        domega_b=other_state.omega_b - self.omega_b,
        ddcm_g2b=geometry.DcmToAxis(other_state.dcm_g2b * self.dcm_g2b.T),
        dwing_vel_g=other_state.wing_vel_g - self.wing_vel_g,
        dwing_pos_g=other_state.wing_pos_g - self.wing_pos_g)

  @type_util.RequireMatrixArguments(None, (3, 1))
  def CalcAerodynamicAngles(self, wind_g):
    """Calculates (v_rel, alpha, beta) from the current wing state.

    Args:
      wind_g: A 3-by-1 matrix storing the wind in g coordinates.

    Returns:
      A tuple (v_rel, alpha, beta).
    """
    return geometry.VelocitiesToAerodynamicAngles(
        self.dcm_g2b, self.wing_vel_g, wind_g)

@type_util.RequireMatrixArguments((3, 1), (3, 2), None, None)
def _CalcBridleKnotPos(tether_force_b, bridle_pos, bridle_y_offset,
                       bridle_radius):
  """Calculate the bridle knot position in body coordinates."""
  if np.linalg.norm(tether_force_b) == 0.0:
    tether_force_b = np.matrix([[0.0], [0.0], [1.0]])

  # Calculate the knot point location.  Here we use a bridle
  # coordinate system with its origin at the bridle pivot, its
  # y-axis pointing toward the starboard bridle point and its z-axis
  # pointed at the knot.
  bridle_coord_y = bridle_pos[:, 1] - bridle_pos[:, 0]
  bridle_coord_y /= np.linalg.norm(bridle_coord_y)

  bridle_coord_z = copy.copy(tether_force_b)
  bridle_coord_z -= bridle_coord_y * (np.transpose(bridle_coord_y)
                                      * tether_force_b)
  bridle_coord_z /= np.linalg.norm(bridle_coord_z)

  bridle_coord_origin = (bridle_pos[:, 1] + bridle_pos[:, 0]) * 0.5
  bridle_coord_origin[1] += bridle_y_offset

  return bridle_coord_origin + bridle_coord_z * bridle_radius

class MotorModel(object):

  # pylint: disable=unused-argument
  def CalcMotorForceMomentPos(self, v_rel, alpha, beta, omega_b,
                              thrust, motor_moment_r):
    raise NotImplementedError()

class PureForceMomentMotorModel(MotorModel):

  def __init__(self, rotor_params, pos_com_b):
    self._dcm_r2b = geometry.AngleToDcm(
        0.0, np.arctan2(rotor_params[0]['axis'][2],
                        rotor_params[0]['axis'][0]), 0.0)
    self._pos_com_b = np.matrix(pos_com_b).T

  # pylint: disable=unused-argument
  @type_util.RequireMatrixArguments(None, None, None, None, (3, 1), (1, 1),
                                    (3, 1))
  def CalcMotorForceMomentPos(self, v_rel, alpha, beta, omega_b,
                              thrust, motor_moment_r):
    # NOTE: This neglects motor reaction torques, and assumes that
    # MixRotors cancels the non-zero torque about the center-of-mass
    # that results from pure thrusting.
    motor_force_r = np.matrix([[thrust[0, 0]], [0.0], [0.0]])
    return ForceMomentPos(
        self._dcm_r2b * motor_force_r, self._dcm_r2b * motor_moment_r,
        self._pos_com_b)

class MotorMixerMotorModel(MotorModel):
  """Model the commanded thrust and moment by calling MixRotors."""

  def __init__(self, rotor_databases, air_density, weights, rotor_params,
               rotor_control_params, hover_flight_mode=False):
    self._dcm_r2b = geometry.AngleToDcm(
        0.0, np.arctan2(rotor_params[0]['axis'][2],
                        rotor_params[0]['axis'][0]), 0.0)
    self._rotor_databases = rotor_databases
    self._air_density = air_density
    self._weights = weights
    self._rotor_params = rotor_params
    self._rotor_control_params = rotor_control_params
    self._hover_flight_mode = hover_flight_mode

  @type_util.RequireMatrixArguments(None, None, (3, 1), (1, 1), (3, 1))
  def CalcRotorSpeeds(self, v_rel, omega_b, thrust, motor_moment_r):
    thrust_moment = {
        'thrust': thrust[0, 0],
        'moment': [motor_moment_r[i, 0] for i in range(3)]
    }

    return actuator_util.MixRotors(
        thrust_moment, self._weights, v_rel, [omega_b[i, 0] for i in range(3)],
        control_types.kStackingStateNormal, self._hover_flight_mode,
        self._air_density, self._rotor_params, self._rotor_control_params)

  @type_util.RequireMatrixArguments(None, None, None, None, (3, 1), (1, 1),
                                    (3, 1))
  def CalcMotorForceMomentPos(self, v_rel, alpha, beta, omega_b,
                              thrust, motor_moment_r):
    rotor_speeds = self.CalcRotorSpeeds(v_rel, omega_b, thrust, motor_moment_r)
    total_force = np.matrix(np.zeros((3, 1)))
    total_moment = np.matrix(np.zeros((3, 1)))

    v_rel_b = geometry.AerodynamicAnglesToRelativeVelocity(v_rel, alpha, beta)
    for i in range(rotor_speeds.shape[0]):
      rotor_speed = rotor_speeds[i, 0]
      if self._rotor_params[i]['dir'] == system_types.kPositiveX:
        direction = 1.0
      else:
        direction = -1.0
      rotor_velocity = direction * rotor_speed

      rotor_pos_b = np.matrix(self._rotor_params[i]['pos']).T
      v_freestream = np.dot(
          self._dcm_r2b[:, 0].T, v_rel_b + np.cross(omega_b.T, rotor_pos_b.T).T)
      v_freestream *= (1.0 - self._rotor_params[i]['local_pressure_coeff'])**0.5

      rotor_thrust = self._rotor_databases[i].CalcThrust(
          rotor_speed, v_freestream[0, 0], self._air_density)
      rotor_torque = direction * self._rotor_databases[i].CalcTorque(
          rotor_speed, v_freestream[0, 0], self._air_density)

      motor_force_b = self._dcm_r2b * np.matrix([[rotor_thrust], [0.0], [0.0]])

      lever_arm_moment_b = np.cross(
          rotor_pos_b.T, motor_force_b.T).T
      aero_moment_b = self._dcm_r2b * np.matrix([[rotor_torque], [0.0], [0.0]])
      gyro_moment_b = np.cross(
          self._rotor_params[i]['I'] * rotor_velocity * self._dcm_r2b[:, 0].T,
          omega_b.T).T

      total_force += motor_force_b
      total_moment += lever_arm_moment_b + aero_moment_b + gyro_moment_b

    return ForceMomentPos(
        total_force, total_moment, np.matrix(np.zeros((3, 1))))

class TetherForceModel(object):

  # pylint: disable=unused-argument
  def CalcBodyForce(self, dcm_g2b, wing_pos_g, wing_vel_g, wind_g):
    raise NotImplementedError()

class ConstantTetherForceModel(TetherForceModel):
  """Simple model of tether force as constant in ground coordinates."""

  def __init__(self, force_g):
    self.SetForce(force_g)

  # pylint: disable=unused-argument
  @type_util.RequireMatrixArguments(None, (3, 3), (3, 1), (3, 1), (3, 1))
  def CalcBodyForce(self, dcm_g2b, wing_pos_g, wing_vel_g, wind_g):
    """Calculate the tether force in body coordinates.

    Args:
      dcm_g2b: DCM rotating ground to body coordinates (3-by-3 np.matrix).
      wing_pos_g: Wing positon [m] in ground coordinates (unused
          3-by-1 np.matrix).
      wing_vel_g: Wing velocity [m/s] in ground coordinates (unused
          3-by-1 np.matrix).
      wind_g: Wind velocity [m/s] in ground coordinates (unused 3-by-1
          np.matrix).

    Returns:
      A 3-by-1 np.matrix storing the tether force in body coordinates.
    """
    return dcm_g2b * self._force_g

  @type_util.RequireMatrixArguments(None, (3, 1))
  def SetForce(self, force_g):
    """Update the force vector.

    Args:
      force_g: New tether force in ground coordinates.
    """
    self._force_g = copy.copy(force_g)

class SimpleSpringTetherForceModel(TetherForceModel):
  """Model of tether force as a simple spring, including bridle interactions."""

  def __init__(self, spring_const, system_params):
    tether_params = system_params['tether']
    wing_params = system_params['wing']

    self._spring_const = spring_const
    self._tether_length = tether_params['length']
    self._tether_drag_area = (0.25 * tether_params['section_drag_coeff']
                              * tether_params['length']
                              * tether_params['outer_diameter'])
    self._air_density = system_params['phys']['rho']
    self._bridle_pos = np.matrix(wing_params['bridle_pos']).T
    self._bridle_y_offset = wing_params['bridle_y_offset']
    self._bridle_radius = wing_params['bridle_rad']

  # pylint: disable=unused-argument
  @type_util.RequireMatrixArguments(None, (3, 3), (3, 1), (3, 1), (3, 1))
  def CalcBodyForce(self, dcm_g2b, wing_pos_g, wing_vel_g, wind_g):
    """Calculate the tether force in body coordinates.

    Args:
      dcm_g2b: DCM rotating ground to body coordinates (3-by-3 np.matrix).
      wing_pos_g: Wing positon [m] in ground coordinates (3-by-1 np.matrix).
      wing_vel_g: Wing velocity [m/s] in ground coordinates (3-by-1 np.matrix).
      wind_g: Wind velocity [m/s] in ground coordinates (3-by-1 np.matrix).

    Returns:
      A 3-by-1 np.matrix storing the tether force in body coordinates.
    """
    # This intentionally ignores the small offset from the GSG
    # position for simplicity.
    bridle_knot_b = _CalcBridleKnotPos(dcm_g2b * -wing_pos_g,
                                       self._bridle_pos,
                                       self._bridle_y_offset,
                                       self._bridle_radius)
    bridle_knot_g = wing_pos_g + dcm_g2b.T * bridle_knot_b
    tension = self._spring_const * (np.linalg.norm(bridle_knot_g)
                                    - self._tether_length)
    spring_force_g = -tension * bridle_knot_g / np.linalg.norm(bridle_knot_g)

    airspeed = np.linalg.norm(wing_vel_g - wind_g)
    drag = 0.5 * self._air_density * airspeed**2.0 * self._tether_drag_area
    drag_force_g = drag * (wind_g - wing_vel_g) / max(airspeed, 0.1)
    return dcm_g2b * (spring_force_g + drag_force_g)

class CatenaryTetherForceModel(TetherForceModel):
  """Model of tether force using catenary tension and rigid-rod drag."""

  def __init__(self, tether_params, gsg_pos_g, bridle_radius, g, air_density):
    """Create a catenary tether force model.

    Args:
      tether_params: TetherParams dictionary.
      gsg_pos_g: Position [m] of the GSG in the g-frame.
      bridle_radius: Bridle radius [m] of the kite.
      g: Gravitational acceleration [m/s^2].
      air_density: Air density [kg/m^3].
    """
    self._gsg_pos_g = np.matrix(np.reshape(gsg_pos_g, (3, 1)))
    self._length = tether_params['length'] + bridle_radius
    self._weight = tether_params['length'] * tether_params['linear_density'] * g
    self._section_drag_coeff = tether_params['section_drag_coeff']
    self._outer_diameter = tether_params['outer_diameter']
    self._air_density = air_density

  # pylint: disable=unused-argument
  @type_util.RequireMatrixArguments(None, (3, 3), (3, 1), (3, 1), (3, 1))
  def CalcBodyForce(self, dcm_g2b, wing_pos_g, wing_vel_g, wind_g):
    """Calculate the tether force in body coordinates.

    Args:
      dcm_g2b: DCM rotating ground to body coordinates (3-by-3 np.matrix).
      wing_pos_g: Wing positon [m] in ground coordinates (3-by-1 np.matrix).
      wing_vel_g: Wing velocity [m/s] in ground coordinates (3-by-1 np.matrix).
      wind_g: Wind velocity [m/s] in ground coordinates (unused 3-by-1
          np.matrix).

    Returns:
      A 3-by-1 np.matrix storing the tether force in body coordinates.
    """
    # Calculate catenary tension.
    horizontal_distance = (wing_pos_g[0, 0]**2.0 + wing_pos_g[1, 0]**2.0)**0.5
    vertical_distance = self._gsg_pos_g[2, 0] - wing_pos_g[2, 0]
    (h, v) = catenary.DimensionlessTensionsFromPoint(
        horizontal_distance / self._length,
        vertical_distance / self._length)
    azi = np.arctan2(wing_pos_g[1, 0], wing_pos_g[0, 0])
    tension_g = self._weight * np.matrix(
        [[-h * np.cos(azi)], [-h * np.sin(azi)], [v]])

    # Calculate drag reaction force on the wing. This is calculated by modeling
    # the tether as a rigid rod that is pinned at the GSG and rotating at fixed
    # angular velocity.
    #
    # Let
    #     CD  = cross-sectional drag coefficient
    #     s   = diameter of the rod
    #     L   = length of rod
    #     V   = velocity of the free end of the rod
    #     rho = air density
    # The drag dD along a segment of the rod with length dx at distance x from
    # the fixed end is
    #     dD(x) = 1/2 * rho * v(x)^2 * CD * s * dx.
    # Therefore,
    #     dD/dx = 1/2 * rho * v(x)^2 * CD * s.
    # The velocity of the segment is v(x) = x/L * V, so
    #     dD/dx = 1/2 * rho * x^2 / L^2 * V^2 * CD * s
    # From this, we obtain the differential moment about the fixed end:
    #     dM/dx = x * dD/dx = 1/2 * rho * x^3 / L^2 * V^2 * CD * s.
    # Integrating from x=0 to x=L yields the total moment due to drag,
    #     M = 1/8 * rho * L^2 * V^2 * CD * s.
    # Force at the fixed end induces no moment, so the drag moment must be
    # entirely balanced by a reaction force at the free end (i.e. the kite).
    # The magnitude of this force, R, is
    #     R = M / L = 1/8 * rho * L * V^2 * CD * s.
    #
    # Here, we treat the rod as extending from the GSG to the body frame origin,
    # and we use the wing velocity normal to the rod to determine V.
    gsg_to_wing_g = wing_pos_g - self._gsg_pos_g
    gsg_to_wing_dir_g = gsg_to_wing_g / np.linalg.norm(gsg_to_wing_g)

    normal_vel_g = (wing_vel_g
                    - float(wing_vel_g.T * gsg_to_wing_dir_g)
                    * gsg_to_wing_dir_g)
    normal_vel_mag = np.linalg.norm(normal_vel_g)
    drag_direction_g = -normal_vel_g / normal_vel_mag
    drag_g = (1.0 / 8.0 * self._air_density * np.linalg.norm(gsg_to_wing_g)
              * normal_vel_mag**2.0 * self._section_drag_coeff
              * self._outer_diameter * drag_direction_g)

    return dcm_g2b * (tension_g + drag_g)

class SwigAeroModel(object):
  """Swig import of the simulator aerodynamics model."""

  def __init__(self):
    self._aero = physics.Aero(physics.GetAeroSimParams())

  @type_util.RequireMatrixArguments(None, None, None, None,
                                    (system_types.kNumFlaps, 1), (3, 1),
                                    None)
  def CalcFMCoeff(self, alpha, beta, reynolds_number, flaps, omega_hat,
                  thrust_coeff):
    """Calculates force and moment coefficients from the Swig database."""
    omega_hat_vec3 = physics.Vec3()
    omega_hat_vec3.x = omega_hat[0, 0]
    omega_hat_vec3.y = omega_hat[1, 0]
    omega_hat_vec3.z = omega_hat[2, 0]
    flaps_vec = physics.VecWrapper(system_types.kNumFlaps)
    for i in range(system_types.kNumFlaps):
      flaps_vec.SetValue(i, flaps[i, 0])
    force_moment = physics.ForceMoment()
    self._aero.CalcForceMomentCoeff(alpha, beta, omega_hat_vec3.this,
                                    flaps_vec.GetVec(), reynolds_number,
                                    force_moment.this, thrust_coeff)
    force_moment_coeff = (np.matrix([[force_moment.force.x],
                                     [force_moment.force.y],
                                     [force_moment.force.z]]),
                          np.matrix([[force_moment.moment.x],
                                     [force_moment.moment.y],
                                     [force_moment.moment.z]]))
    return force_moment_coeff

class Wing(object):
  """Simplified model of the wing for control design.

  The Wing class stores parameters defined by the environment (air
  density, gravitational constant), a stateless tether force model,
  a stateless aerodynamic model, and a nominal orientation.

  It provides functions for calculating the ODEs that govern a 6-DOF
  rigid body model.
  """

  def __init__(self, system_params, sim_params, aero_model, motor_model,
               tether_force_model):
    """Constructs a Wing model.

    Args:
      system_params: A system parameters structure from mconfig.
      sim_params: A simulator parameters structure from mconfig.
      aero_model: A Python class implementing a function CalcFMCoeff.  See
          SwigAeroModel in this module as an example.
      motor_model: A MotorModel.
      tether_force_model: A TetherForceModel.
    """
    self._wing_area = system_params['wing']['A']
    self._wing_span = system_params['wing']['b']
    self._wing_chord = system_params['wing']['c']
    self._wing_mass = system_params['wing']['m']
    self._wing_inertia_matrix = np.matrix(system_params['wing']['I']['d'])
    self._pos_com_b = np.matrix(system_params['wing']['center_of_mass_pos']).T

    # Bridle parameters.
    self._bridle_pos = np.matrix(system_params['wing']['bridle_pos']).T
    self._bridle_y_offset = system_params['wing']['bridle_y_offset']
    self._bridle_radius = system_params['wing']['bridle_rad']

    # Physics parameters.
    self._g_g = np.matrix([[0.0], [0.0], [system_params['phys']['g']]])
    self._air_density = system_params['phys']['rho']
    self._dynamic_viscosity = sim_params['phys_sim']['dynamic_viscosity']

    self._aero_model = aero_model
    self._motor_model = motor_model
    self._tether_force_model = tether_force_model

  @type_util.RequireMatrixArguments(None, (3, 3))
  def _CalcGravityForceMomentPos(self, dcm_g2b):
    return ForceMomentPos(dcm_g2b * (self._wing_mass * self._g_g),
                          np.matrix(np.zeros((3, 1))), self._pos_com_b)

  @type_util.RequireMatrixArguments(None, None, None, None, (3, 1), (1, 1),
                                    (3, 1))
  def _CalcMotorForceMomentPos(self, v_rel, alpha, beta, omega_b,
                               thrust, motor_moment):
    """Calculates the motor forces and moments."""
    return self._motor_model.CalcMotorForceMomentPos(
        v_rel, alpha, beta, omega_b, thrust, motor_moment)

  @type_util.RequireMatrixArguments(None, (3, 3), (3, 1), (3, 1), (3, 1))
  def _CalcTetherForceMomentPos(self, dcm_g2b, wing_pos_g, wing_vel_g, wind_g):
    tether_force_b = self._tether_force_model.CalcBodyForce(dcm_g2b, wing_pos_g,
                                                            wing_vel_g, wind_g)
    return ForceMomentPos(tether_force_b, np.matrix(np.zeros((3, 1))),
                          _CalcBridleKnotPos(tether_force_b, self._bridle_pos,
                                             self._bridle_y_offset,
                                             self._bridle_radius))

  def CalcTetherForceG(self, state, inputs):
    return state.dcm_g2b.T * self._tether_force_model.CalcBodyForce(
        state.dcm_g2b, state.wing_pos_g, state.wing_vel_g, inputs.wind_g)

  def CalcTetherTensionRollPitch(self, state, inputs):
    tether_force_b = self._tether_force_model.CalcBodyForce(
        state.dcm_g2b, state.wing_pos_g, state.wing_vel_g, inputs.wind_g)
    return geometry.TetherForceCartToSph(tether_force_b)

  @type_util.RequireMatrixArguments(None, None, None, None, (3, 1),
                                    (system_types.kNumFlaps, 1), None)
  def CalcAeroForceMomentPos(self, v_rel, alpha, beta, omega_b, flaps,
                             thrust_coeff):
    """Calculate the aerodynamic force and moments on the wing.

    Args:
      v_rel: Airspeed [m/s].
      alpha: Angle-of-attack [rad].
      beta: Angle-of-sideslip [rad].
      omega_b: Wing body-rates [rad/s] (3-by-1 np.matrix).
      flaps: Flap deflections (kNumFlaps-by-1 np.matrix).
      thrust_coeff: Thrust coefficient [#] using wind turbine convention.

    Returns:
      (ForceMomentPos in body coordinates, force coeffs., moment coeffs.)
    """
    reynolds_number = ((v_rel * self._wing_chord * self._air_density)
                       / self._dynamic_viscosity)
    dynamic_pressure = 0.5 * self._air_density * v_rel**2.0

    length_scale = np.matrix([[self._wing_span],
                              [self._wing_chord],
                              [self._wing_span]])
    omega_hat = np.multiply(omega_b, length_scale) / (2.0 * v_rel)
    (cf, cm) = self._aero_model.CalcFMCoeff(alpha, beta, reynolds_number,
                                            flaps, omega_hat, thrust_coeff)

    return (ForceMomentPos(dynamic_pressure * self._wing_area * cf,
                           (dynamic_pressure * self._wing_area
                            * np.multiply(length_scale, cm)),
                           np.matrix(np.zeros((3, 1)))),
            cf, cm)

  def _BodyForceMomentPosToComForceMoment(self, force_moment_pos_list):
    force = np.matrix(np.zeros((3, 1)))
    moment = np.matrix(np.zeros((3, 1)))
    for force_moment_pos in force_moment_pos_list:
      force += force_moment_pos.force
      moment += force_moment_pos.moment
      moment += np.cross(force_moment_pos.pos - self._pos_com_b,
                         force_moment_pos.force, axis=0)
    return ForceMoment(force, moment)

  def CalcDeriv(self, state, inputs):
    """Calculates the derivative of the wing state vector.

    Args:
      state: A WingState.
      inputs: A WingInputs.

    Returns:
      A WingState.Tangent containing the derivative of the state.
    """
    euler_moment = np.cross(self._wing_inertia_matrix * state.omega_b,
                            state.omega_b, axis=0)

    v_rel, alpha, beta = state.CalcAerodynamicAngles(inputs.wind_g)

    # Fixing total thrust coefficient to 0.0 for this application.
    # NOTE: By accounting for the rotor wake effect on the tail,
    # we found that the synthesized gains yield worse flight quality than when
    # the effect is ignored (see b/110491871 for details).
    thrust_coeff = 0.0
    aero_force_moment_pos, _, _ = self.CalcAeroForceMomentPos(
        v_rel, alpha, beta, state.omega_b, inputs.flaps, thrust_coeff)

    force_moment_com = self._BodyForceMomentPosToComForceMoment([
        self._CalcGravityForceMomentPos(state.dcm_g2b),
        self._CalcMotorForceMomentPos(
            v_rel, alpha, beta, state.omega_b, inputs.thrust,
            inputs.motor_moment),
        self._CalcTetherForceMomentPos(state.dcm_g2b, state.wing_pos_g,
                                       state.wing_vel_g, inputs.wind_g),
        aero_force_moment_pos,
        ForceMomentPos(np.matrix(np.zeros((3, 1))), euler_moment,
                       np.matrix(np.zeros((3, 1))))
    ])

    # Calculate center-of-mass acceleration.
    accel_com_g = (state.dcm_g2b.T * force_moment_com.force) / self._wing_mass

    # Calculate body angular acceleration.
    omega_b_dot = np.matrix(np.linalg.solve(self._wing_inertia_matrix,
                                            force_moment_com.moment))

    wing_accel_g = accel_com_g - state.dcm_g2b.T * (
        np.cross(state.omega_b,
                 np.cross(state.omega_b, self._pos_com_b, axis=0), axis=0)
        + np.cross(omega_b_dot, self._pos_com_b, axis=0))

    return WingState.Tangent(domega_b=omega_b_dot, ddcm_g2b=state.omega_b,
                             dwing_vel_g=wing_accel_g,
                             dwing_pos_g=state.wing_vel_g)

  def CalcDVbCom(self, state, state_dot):
    """Calculates the rate of change of Vb for unit tests."""
    return (state.dcm_g2b * state_dot.dwing_vel_g
            - np.cross(state.omega_b, state.dcm_g2b * state.wing_vel_g, axis=0)
            + np.cross(state_dot.domega_b, self._pos_com_b, axis=0))

  def CalcEnergy(self, state):
    """Calculates energy of the rigid body model for unit tests."""
    wing_com_pos_g = state.wing_pos_g + state.dcm_g2b.T * self._pos_com_b
    wing_com_vel_g = (state.wing_vel_g
                      + (state.dcm_g2b.T
                         * np.cross(state.omega_b, self._pos_com_b, axis=0)))
    return ((0.5 * np.transpose(state.omega_b)
             * self._wing_inertia_matrix * state.omega_b)
            + (0.5 * self._wing_mass * np.transpose(wing_com_vel_g)
               * wing_com_vel_g)
            - self._wing_mass * np.transpose(self._g_g) * wing_com_pos_g)[0, 0]

def CalcLinearization(f, state, inputs, state_step_sizes, input_step_sizes):
  """Calculate the system matrices for the Wing model.

  Produces a linearized model:

    f(x + dx, u + du) ~ f(x) + A * dx + B * du

  where f is an arbitrary function, x is the wing state and u are
  the wing inputs.

  Args:
    f: A function mapping an n-by-1 np.matrix and an m-by-1 np.matrix to
        a n-by-1 np.matrix.
    state: An instance of a state class from type_util.
    inputs: An instance of a named vector from type_util.
    state_step_sizes: A vector of step sizes for the state.
    input_step_sizes: A vector of step sizes for the inputs.

  Returns:
    A tuple (A, B) where A and B are both of type np.matrix.
  """
  num_states = state.Tangent.GetDim()
  num_inputs = inputs.GetDim()
  num_outputs = f(state, inputs).shape[0]

  dfdx = np.matrix(np.zeros((num_outputs, num_states)))
  dfdu = np.matrix(np.zeros((num_outputs, num_inputs)))
  for i in range(num_states):
    h = state_step_sizes[i, 0]
    e = state.Tangent.FromVector(np.matrix([
        [1.0 if j == i else 0.0] for j in range(num_states)]))

    dfdx[:, i] = (f(state.Increment(e, step=h), inputs)
                  - f(state.Increment(e, step=-h), inputs)) / (2.0 * h)

  for i in range(num_inputs):
    h = input_step_sizes[i, 0]
    e = np.matrix([[1.0 if j == i else 0.0] for j in range(num_inputs)])
    dfdu[:, i] = (
        f(state, inputs.FromVector(inputs.ToVector() + h * e))
        - f(state, inputs.FromVector(inputs.ToVector() - h * e))) / (2.0 * h)

  return (dfdx, dfdu)