"""
因子数据处理时常用的一些函数
1. 上市时间过滤，剔除新股，默认剔除上市时间 200 天以内的股票, 暂时采用聚宽接口
2. 极值处理, 支持 MAD, STD, Quantile 方式
3. 标准化处理，暂时支持 std
4. 获取行业数据，暂时采用聚宽接口，支持 "sw_l1", "sw_l2", "sw_l3", "jq_l1", "jq_l2", "zjw"
5. 获取市值数据，使用本地接口，支持 'avg', 'mktcap', 'sqrt_mktcap', 'ln_mktcap',
    'cmktcap', 'sqrt_cmktcap', 'ln_cmktcap'
6. 行业中性化与市值中性化
"""
from functools import partial
from typing import List, Tuple, Union

try:
    import jqdatasdk
except:
    print("jqdatasdk not installed")
import numpy as np
import pandas as pd

import statsmodels.api as sm
from QUANTAXIS.QAAnalysis.QAAnalysis_block import QAAnalysis_block
from QUANTAXIS.QAFactor.utils import QA_fmt_code_list
from QUANTAXIS.QAFactor.fetcher import (QA_fetch_stock_basic, QA_fetch_industry_adv)

def QA_fmt_factor(factor: Union[pd.Series, pd.DataFrame]):
    """
    将 factor 格式化
    """
    # 索引设置
    if isinstance(factor, pd.DataFrame):
        factor = factor.stack(dropna=False)
    level_0 = 'datetime'
    try:
        pd.Timestamp(factor.index.levels[0][0])
    except:
        try:
            level_0 = 'code'
            pd.Timestamp(factor.index.levels[1][0])
        except:
            raise ValueError("错误，索引格式不正确")
    if level_0 == 'code':
        factor = factor.swaplevel()
    factor.index.names = ["datetime", "code"]
    factor.index = factor.index.map(lambda x: (pd.Timestamp(x[0]), x[1]))
    return factor.sort_index()

def QA_winsorize_factor(
        factor: Union[pd.Series,
                      pd.DataFrame],
        grouper: list = ["datetime"],
        extreme_scale: float = 1.4826,
        extreme_method: str = "mad",
        extreme_inclusive: bool = False,
        quant: list = None,
        inf2nan: bool = True,
        nan2mean: bool = True,
) -> pd.Series:
    """
    对输入的 Series 进行极值处理
    :param factor: 因子数据, ['日期' '资产'] 的 MultiIndex 或者 ['日期'] 为索引，['资产'] 为列索引的透视表
    :param grouper: 分组数据
    :param extreme_scale: 极值处理对应的 scale
    :param extreme_method: 极值处理方式
    :param extreme_inclusive: 是否保留边界外极值为边界值，默认不保留
    :param inf2nan: 是否将 np.inf -np.inf 转换为 np.nan
    :param quant: 0 ~ 1 之间的范围
    :param nan2mean: 是否将 np.nan 替换为行业平均值
    """

    def _mad_cut(x, scale, inclusive, nan2mean):
        diff = (x - x.median()).apply(abs)
        mad = diff.median()
        upper_limit = x.median() + scale * mad
        lower_limit = x.median() - scale * mad
        if inclusive:
            x.loc[x < lower_limit] = lower_limit
            x.loc[x > upper_limit] = upper_limit
            x.clip(lower_limit, upper_limit, inplace=True)
        else:
            x = x.loc[(x >= lower_limit) & (x <= upper_limit)]
        if nan2mean:
            return x.fillna(np.nanmean(x))
        return x

    def _std_cut(x, scale, inclusive, nan2mean):
        std = x.std()
        upper_limit = x.mean() + scale * std
        lower_limit = x.mean() - scale * std
        if inclusive:
            x.loc[x < lower_limit] = lower_limit
            x.loc[x > upper_limit] = upper_limit
            factor.clip(lower_limit, upper_limit, inplace=True)
        else:
            x = x.loc[(x > lower_limit) & (x < upper_limit)]
        if nan2mean:
            return x.fillna(np.nanmean(x))
        else:
            return x

    def _quant_cut(x, quant, inclusive, nan2mean):
        if inclusive:
            x.loc[x < quant[0]] = quant[0]
            x.loc[x > quant[1]] = quant[1]
            x = np.quantile(x, quant, interpolation="nearest")
        else:
            x = np.quantile(x, quant, interpolation="nearest")
        if nan2mean:
            return x.fillna(np.nanmean(x))
        else:
            return x

    if extreme_method not in ["mad", "std", "quant"]:
        raise ValueError("参数为 mad/std/quant, 仅支持这极值处理方式")
    # 对因子进行格式化
    factor = QA_fmt_factor(factor)
    # inf2nan
    if inf2nan:
        factor.replace([np.inf, -np.inf], np.nan, inplace=True)
    # 极值处理
    if extreme_method == "mad":
        factor = (
            factor.groupby(grouper).apply(
                lambda x:
                _mad_cut(x,
                         extreme_scale,
                         extreme_inclusive,
                         nan2mean)
            ).droplevel(level=1)
        )
        return factor
    elif extreme_method == "std":
        factor = (
            factor.groupby(grouper).apply(
                lambda x:
                _std_cut(x,
                         extreme_scale,
                         extreme_inclusive,
                         nan2mean)
            ).droplevel(level=1)
        )
        return factor
    elif extreme_method == "quant":
        if quant is None:
            raise ValueError("")
        factor = (
            factor.groupby(grouper)
            .apply(lambda x: _quant_cut(x,
                                        quant,
                                        extreme_inclusive,
                                        nan2mean)).droplevel(level=1)
        )
        return factor

def QA_standardize_factor(
        factor: pd.Series,
        standard_method: str = "normal",
        inf2nan: bool = True
) -> pd.Series:
    """
    对输入的因子值进行标准化处理，默认为 Z-score 处理, 其他处理方式待实现
    :param data: 输入的 DataFrame, 索引为股票代码
    :param standard_method: 标准化处理方法，目前接受 "normal"/"weighted"/"john"
    :param inf2nan: 是否将 np.inf 转换为 np.nan
    """
    factor = QA_fmt_factor(factor)
    if standard_method not in ["normal", "weighted", "john"]:
        raise ValueError("仅接受 normal/weighted/john 参数")
    if standard_method == "normal":
        std = factor.std()
        mean = factor.mean()
        factor = (factor - mean) / std
    elif standard_method == "weighted":
        # TODO: 市值加权标准化方法待实现
        raise ValueError("市值加权标准化方法待实现")
    else:
        # TODO: Johnson 变化法
        raise ValueError("John 变化法待实现")
    return factor

def QA_fetch_get_factor_groupby(
        factor: pd.Series,
        industry_cls: str = "sw_l1",
        detailed: bool = False
) -> pd.DataFrame:
    """
    获取因子的行业暴露, 注意，返回的值是 pd.DataFrame 格式，包含原因子值，附加一列
    因子对应的行业信息 (需先自行导入聚宽本地 sdk 并登陆)

    参数
    ---
    :param factor: 因子值，索引为 ['日期' '资产']
    :param industry_cls: 行业分类，默认为申万 1 级行业
    :param detailed: 是否使用详细模式，默认为 False, 即取因子日期最后一日的行业信息

    返回值
    ---
    :return: 因子数据, 包括因子值，因子对应行业
    """
    warnings.warn("请先自行导入聚宽本地 sdk 并登陆", UserWarning)
    # 因子格式化
    factor = QA_fmt_factor(factor)
    merged_data = pd.DataFrame(factor.copy().rename("factor"))
    # 股票代码格式化
    stock_list = QA_fmt_code_list(
        factor.index.get_level_values("code").drop_duplicates(),
    )
    # 非详细模式， 行业数据采用当前日期
    ss = pd.Series()
    if detailed:
        # start_time = str(min(factor.index.get_level_values("datetime")))[:10]
        # end_time = str(max(factor.index.get_level_values("datetime")))[:10]
        # date_range = list(
        #     map(pd.Timestamp, QA_util_get_trade_range(start_time, end_time))
        # )
        date_range = (
            factor.index.get_level_values(
                "datetime").drop_duplicates().tolist()
        )

        industry = pd.DataFrame()
        for cursor_date in date_range:
            df_tmp = QA_fetch_industry_adv(code = stock_list, cursor_date = cursor_date)[["code", "industry_name"]]
            df_tmp["date"] = cursor_date
            industry = industry.append(df_tmp)
        ss = industry.set_index(["date", "code"])["industry_name"]
        # industries = map(
        #     partial(jqdatasdk.get_industry,
        #             stock_list),
        #     date_range
        # )
        # industries = {
        #     d: {
        #         s: ind.get(s).get(industry_cls,
        #                           dict()).get("industry_name",
        #                                       "NA")
        #         for s in stock_list
        #     }
        #     for d,
        #     ind in zip(date_range,
        #                industries)
        # }
    else:
        end_time = str(max(factor.index.get_level_values("datetime")))[:10]
        date_range = [pd.Timestamp(end_time)]
        # industries = jqdatasdk.get_industry(stock_list, end_time)
        ss = QA_fetch_industry_adv(stock_list, end_time)[["code", "industry_name"]].set_index(["date", "code"])["industry_name"]
        # industries = {
        #     d: {
        #         s: industries.get(s).get(industry_cls,
        #                                  dict()).get("industry_name",
        #                                              "NA")
        #         for s in stock_list
        #     }
        #     for d in date_range
        # }
    # 可能历史上股票没有行业信息，用之后的行业信息往前填充
    merged_data["date"] = merged_data.index.get_level_values("datetime").map(
        lambda x: x.date()
    )
    merged_data = (
        merged_data.reset_index().set_index(
            ["date",
             "code"]
        ).assign(group=ss).reset_index().set_index(["datetime",
                                                          "code"]
                                                         ).drop("date",
                                                                axis=1)
    )
    group = merged_data["group"].unstack().bfill().stack()
    merged_data["group"] = group
    return merged_data

def QA_fetch_factor_weight(
        factor: pd.Series,
        weight_cls: str = "mktcap",
        detailed: bool = True
) -> pd.DataFrame:
    """
    获取因子的市值暴露, 注意，返回的值是 pd.DataFrame 格式，包含原因子值，附加一列
    因子对应的加权信息

    参数
    ---
    :param factor: 因子值，索引为 ['日期' '资产']
    :param weight_cls: 权重信息，默认加权方式为总市值加权
    :param detailed: 默认为 True, 如果为 False, 取因子最后一日的加权信息

    返回值
    ---
    :return: 因子数据, 包括因子值，因子对应行业
    """
    # 因子格式化
    factor = QA_fmt_factor(factor)
    merged_data = pd.DataFrame(factor.copy().rename("factor"))

    # 股票代码格式化
    code_list = factor.index.get_level_values(
        "code").drop_duplicates().tolist()
    # 非详细模式， 加权数据采用当前日期
    if detailed:
        # start_time = str(min(factor.index.get_level_values("datetime")))[:10]
        end_time = str(max(factor.index.get_level_values("datetime")))[:10]
        # date_range = list(
        #     map(pd.Timestamp, QA_util_get_trade_range(start_time, end_time))
        # )
        date_range = (
            factor.index.get_level_values(
                "datetime").drop_duplicates().tolist()
        )

    else:
        date_range = [pd.Timestamp(end_time)]

    if weight_cls == "avg":
        merged_data["weight"] = 1.0
        return merged_data

    df_local = QAAnalysis_block(
        code=code_list,
        start=date_range[0],
        end=date_range[-1]
    ).market_value
    if weight_cls == "mktcap":
        df_local = df_local.reset_index().pivot(
            index="date",
            columns="code",
            values="mv"
        )
    elif weight_cls == "sqrt_mktcap":
        df_local = (
            df_local.reset_index().pivot(
                index="date",
                columns="code",
                values="mv"
            ).transform("sqrt")
        )
    elif weight_cls == "ln_mktcap":
        df_local = (
            df_local.reset_index().pivot(
                index="date",
                columns="code",
                values="mv"
            ).transform("ln")
        )
    elif weight_cls == "cmktcap":
        df_local = df_local.reset_index().pivot(
            index="date",
            columns="code",
            values="liquidity_mv"
        )
    elif weight_cls == "sqrt_cmktcap":
        df_local = (
            df_local.reset_index().pivot(
                index="date",
                columns="code",
                values="liquidity_mv"
            ).transform("sqrt")
        )
    elif weight_cls == "ln_cmktcap":
        df_local = (
            df_local.reset_index().pivot(
                index="date",
                columns="code",
                values="liquidity_mv"
            ).transform("ln")
        )
    else:
        raise ValueError(f"{weight_cls} 加权方式未实现")
    merged_data["date"] = merged_data.index.get_level_values("datetime").map(
        lambda x: x.date()
    )
    merged_data = (
        merged_data.reset_index().set_index(
            ["date",
             "code"]
        ).assign(weight=df_local.stack()
                 ).reset_index().set_index(["datetime",
                                            "code"]).drop("date",
                                                          axis=1)
    )

    weight = merged_data["weight"].unstack().bfill().stack()
    merged_data["weight"] = weight
    return merged_data

def QA_fetch_get_factor_start_date(factor: pd.Series) -> pd.DataFrame:
    """
    获取因子池上市时间, 注意，请自行登陆聚宽本地 sdk

    参数
    ---
    :param factor: 因子值，索引为 ['日期' '资产']

    返回值
    ---
    :return: 因子数据
    """
    # 因子格式化
    factor = QA_fmt_factor(factor.copy())
    merged_data = pd.DataFrame(factor.rename("factor"))
    # 股票代码格式化
    stock_list = QA_fmt_code_list(
        factor.index.get_level_values("code").drop_duplicates(),
        style="jq"
    )
    # 上市时间获取
    df_local = jqdatasdk.get_all_securities(types="stock")
    intersection = df_local.index.intersection(stock_list)
    ss = df_local.loc[intersection]["start_date"]
    ss.index = ss.index.map(lambda x: x[:6])
    # 拼接上市时间
    merged_data = merged_data.loc[(slice(None), list(ss.index)), :]
    merged_data["start_date"] = merged_data.index.map(lambda x: ss.loc[x[1]]
                                                      ).tolist()
    return merged_data

def QA_fetch_factor_start_date(factor: pd.Series) -> pd.DataFrame:
    """
    获取因子池上市时间，本地获取接口，使用前先保存股票基本信息
    """
    factor = QA_fmt_factor(factor.copy())
    merged_data = pd.DataFrame(factor.rename("factor"))
    # 股票代码格式化
    stock_list = QA_fmt_code_list(
        factor.index.get_level_values("code").drop_duplicates()
    )
    # 上市时间获取
    df_local = QA_fetch_stock_basic(status=None).set_index("code")
    intersection = df_local.index.intersection(stock_list)
    ss = df_local.loc[intersection]["list_date"]
    # 拼接上市时间
    merged_data = merged_data.loc[(slice(None), list(ss.index)), :]
    merged_data["start_date"] = merged_data.index.map(lambda x: ss.loc[x[1]]
                                                      ).tolist()
    return merged_data

def QA_neutralize_factor(
        factor: pd.Series,
        weight_cls: str = "avg",
        industry: str = "sw_l1"
) -> pd.DataFrame:
    """
    行业中性处理

    参数
    ---
    :param factor: 因子值
    :param weight_cls: 加权方式
    :param industry_cls: 行业分类方式

    返回值
    ---
    :return: 中性化处理后的因子数据
    """
    if type(mkt_cap) == pd.Series:
        LnMktCap = mkt_cap.apply(lambda x: math.log(x))
        if industry:  # 行业、市值
            dummy_industry = get_industry_exposure(factor, date=date)
            x = pd.concat([LnMktCap, dummy_industry.T], axis=1)
        else:        # 仅市值
            x = LnMktCap
    elif industry:
        dummy_industry = get_industry_exposure(factor, date=date)
        x = dummy_industry.T
    return (sm.OLS(factor[factor], x.astype(float)).fit().resid.rename(factor),)


import numpy as np

# from numba import jit
import torch

from shapely.geometry import Polygon
from pyquaternion import Quaternion

import iou3d

def _hypot(x, y):
    # TODO use torch.hypot for PyTorch >= 1.7
    return torch.sqrt(torch.square(x) + torch.square(y))

"""
Fast (vectorized) boxes API.

Boxes should be numpy.array[B, 7] with (x, y, z, l, w, h, theta)

Boxes follow JRDB base frame convention, x-forward, y-left, z-up, theta is the
angle between x-axis and the length side of bounding boxes
"""

def get_R(boxes):
    """Get rotation matrix R along z axis that transforms a point from box coordinate
    to world coordinate.

    Args:
        boxes (array[B, 7]): (x, y, z, l, w, h, theta) of each box

    Returns:
        Rs (array[B, 3, 3])
    """
    # NOTE plus pi specifically for JRDB, don't know the reason
    theta = boxes[:, 6] + np.pi
    cs, ss = np.cos(theta), np.sin(theta)
    zeros, ones = np.zeros(len(cs)), np.ones(len(cs))
    Rs = np.array(
        [[cs, ss, zeros], [-ss, cs, zeros], [zeros, zeros, ones]], dtype=np.float32
    )  # (3, 3, B)

    return Rs.transpose((2, 0, 1))

def get_R_torch(boxes):
    """Get rotation matrix R along z axis that transforms a point from box coordinate
    to world coordinate.

    Args:
        boxes (tensor[B, 7]): (x, y, z, l, w, h, theta) of each box

    Returns:
        Rs (tensor[B, 3, 3])
    """
    # NOTE plus pi specifically for JRDB, don't know the reason
    theta = boxes[:, 6] + np.pi
    cs, ss = torch.cos(theta), torch.sin(theta)
    Rs = torch.zeros((len(cs), 3, 3), device=boxes.device).float()
    Rs[:, 2, 2] = 1
    Rs[:, 0, 0] = cs
    Rs[:, 0, 1] = ss
    Rs[:, 1, 0] = -ss
    Rs[:, 1, 1] = cs

    return Rs

def boxes_to_corners(boxes, resize_factor=1.0, connect_inds=False):
    """Return xyz coordinates of the eight vertices of the bounding box

    First four points are fl (front left), fr, br, bl on top plane. Last four
    points are same order, but for the bottom plane.

          0 -------- 1        __
         /|         /|        //|
        3 -------- 2 .       //
        | |        | |      front
        . 4 -------- 5
        |/         |/
        7 -------- 6

    To draw a box, do something like

    corners, connect_inds = boxes_to_corners(boxes)
    for corner in corners:
        for inds in connect_inds:
            mlat.plot3d(corner[0, inds], corner[1, inds], corner[2, inds],
                        tube_radius=None, line_width=5)

    Args:
        boxes (array[B, 7]): (x, y, z, l, w, h, theta) of each box
        resize_factor (float): resize box lwh dimension
        connect_inds(bool): true will also return a list of indices for drawing
            the box as line segments

    Returns:
        corners_xyz (array[B, 3, 8])
        connect_inds (tuple[list[int]])
    """
    # in box frame
    c_xyz = np.array(
        [
            [1, 1, -1, -1, 1, 1, -1, -1],
            [-1, 1, 1, -1, -1, 1, 1, -1],
            [1, 1, 1, 1, -1, -1, -1, -1],
        ],
        dtype=np.float32,
    )  # (3, 8)
    c_xyz = 0.5 * c_xyz[np.newaxis, :, :] * boxes[:, 3:6, np.newaxis]  # (B, 3, 8)
    c_xyz = c_xyz * resize_factor

    # to world frame
    R = get_R(boxes)  # (B, 3, 3)
    c_xyz = R @ c_xyz + boxes[:, :3, np.newaxis]  # (B, 3, 8)

    if not connect_inds:
        return c_xyz
    else:
        l1 = [0, 1, 2, 3, 0, 4, 5, 6, 7, 4, 1]
        l2 = [0, 5, 1]
        l3 = [2, 6]
        l4 = [3, 7]
        return c_xyz, (l1, l2, l3, l4)

def boxes_to_corners_torch(boxes):
    """See boxes_to_corners()

    Args:
        boxes (tensor[B, 7]): (x, y, z, l, w, h, theta) of each box

    Returns:
        corners_xyz (tensor[B, 3, 8])
    """
    # in box frame
    c_xyz = torch.tensor(
        [
            [1, 1, -1, -1, 1, 1, -1, -1],
            [-1, 1, 1, -1, -1, 1, 1, -1],
            [1, 1, 1, 1, -1, -1, -1, -1],
        ],
        device=boxes.device,
    ).float()  # (3, 8)
    c_xyz = 0.5 * c_xyz.unsqueeze(dim=0) * boxes[:, 3:6].unsqueeze(dim=2)  # (B, 3, 8)

    # to world frame
    R = get_R_torch(boxes)  # (B, 3, 3)
    c_xyz = R @ c_xyz + boxes[:, :3].unsqueeze(dim=2)  # (B, 3, 8)

    return c_xyz

def boxes_to_central_line(boxes):
    """Return parameter a, b, c such that ax + by + c = 0 is the line that passes
    the box center and is in parallel with the long edge in BEV.

    Args:
        boxes (array[..., 7]): (x, y, z, l, w, h, theta) of each box

    Returns:
        boxes_lp (array[..., 3]): (a, b, c) line parameters of each box
    """
    # in case length is shorter than width
    bmask = boxes[..., 3] < boxes[..., 4]
    theta = -boxes[..., 6]  # not sure why minus is needed
    theta[bmask] -= 0.5 * np.pi

    a = np.tan(theta)
    b = -np.ones_like(a)
    c = -a * boxes[..., 0] - b * boxes[..., 1]
    boxes_lp = np.stack((a, b, c), axis=-1)
    boxes_lp /= np.linalg.norm(boxes_lp, axis=-1, keepdims=True)

    return boxes_lp

def boxes_to_central_line_torch(boxes):
    """See boxes_to_central_line

    Args:
        boxes (tensor[..., 7]): (x, y, z, l, w, h, theta) of each box

    Returns:
        boxes_lp (tensor[..., 3]): (a, b, c) line parameters of each box
    """
    # in case length is shorter than width
    bmask = boxes[..., 3] < boxes[..., 4]
    theta = -boxes[..., 6]  # not sure why minus is needed
    theta[bmask] -= 0.5 * np.pi

    a = torch.tan(theta)
    b = -torch.ones_like(a)
    c = -a * boxes[..., 0] - b * boxes[..., 1]
    boxes_lp = torch.stack((a, b, c), dim=-1)
    boxes_lp /= torch.linalg.norm(boxes_lp, dim=-1, keepdim=True)

    return boxes_lp

def distance_pc_to_boxes_torch(
    pc, boxes, normalize=False, delta_size=0.0, return_in_box_mask=False
):
    """Compute distance between points and matching boxes

    Args:
        pc (tensor[..., 3]): Points
        boxes (tensor[..., 7]): (x, y, z, l, w, h, theta)

    Returns:
        dist (tensor[...])
    """
    # vector from points to boxes, in world frame
    xyz_offs = pc - boxes[..., :3]
    xyz_dist = torch.norm(xyz_offs, dim=-1)  # (...)

    if normalize:
        boxes_diag = 0.5 * torch.norm(boxes[..., 3:6] + delta_size, dim=-1)
        xyz_dist = xyz_dist / boxes_diag

    if not return_in_box_mask:
        return xyz_dist

    # convert offsets to box frame
    out_shape = boxes.shape[:-1]
    R = get_R_torch(boxes.view(-1, 7))  # (B, 3, 3), box to world
    xyz_offs = R.transpose(1, 2) @ xyz_offs.view(-1, 3, 1)  # (B, 3, 1)
    xyz_offs = xyz_offs.view(*out_shape, 3)  # (..., 3)

    hf_lwh = 0.5 * (boxes[..., 3:6] + delta_size)
    ib_mask = torch.logical_and(xyz_offs >= -hf_lwh, xyz_offs <= hf_lwh)
    ib_mask = ib_mask.min(dim=-1)[0]

    return xyz_dist, ib_mask

# @jit(nopython=True)
# def find_closest_boxes(pc, boxes):
#     """For each point, return the index of the box with closest center and the distance

#     Args:
#         pc (array[3, N]): xyz point cloud
#         boxes (array[B, 7]): bounding boxes (x, y, z, l, w, h, theta)

#     Returns:
#         closest_box_inds (array[N,]): The index of the closest box from a point
#         closest_box_dist (array[N,]): Distance to the box center
#     """
#     # offset from box center to points
#     N = pc.shape[1]
#     closest_box_inds = np.empty(N, dtype=np.int32)
#     closest_box_dist = np.empty(N, dtype=np.float32)
#     for n in range(N):
#         pc_n = pc[:, n : n + 1].T
#         pb_dist = pc_n - boxes[:, :3]  # (B, 3)
#         pb_norm_sq = np.square(pb_dist).sum(axis=1)  # (B, )

#         min_idx = pb_norm_sq.argmin()
#         closest_box_inds[n] = min_idx
#         closest_box_dist[n] = np.sqrt(pb_norm_sq[min_idx])

#     return closest_box_inds, closest_box_dist

def find_closest_boxes(pc, boxes):
    """For each point, return the index of the box with closest center and the distance

    Args:
        pc (array[3, N]): xyz point cloud
        boxes (array[B, 7]): bounding boxes (x, y, z, l, w, h, theta)

    Returns:
        closest_box_inds (array[N,]): The index of the closest box from a point
        closest_box_dist (array[N,]): Distance to the box center
    """
    # offset from box center to points
    pc_box_dist = pc[np.newaxis, :, :] - boxes[:, :3, np.newaxis]  # (B, 3, N)
    pc_box_norm = np.linalg.norm(pc_box_dist, axis=1)  # (B, N)
    closest_box_inds = pc_box_norm.argmin(axis=0)  # (N, )
    closest_box_dist = np.take_along_axis(
        pc_box_norm, closest_box_inds[np.newaxis, ...], axis=0
    )[0]

    return closest_box_inds, closest_box_dist

def find_in_box_points(pc, boxes, resize_factor=1.0):
    """For each point, return the indices of the box it belongs to. Only work
    with boxes that is rotated along z (verticle) axis.

    Args:
        pc (array[3, N]): xyz point cloud
        boxes (array[B, 7]): bounding boxes (x, y, z, l, w, h, theta)
        resize_factor (float): resize box dimension

    Returns:
        in_box_mask (array[N,]): True if a point is in at least one box
        closest_box_inds (array[N,]): The index of the closest box from a point
    """
    # offset from box center to points in world coordinate
    pc_box = pc[np.newaxis, :, :] - boxes[:, :3, np.newaxis]  # (B, 3, N)

    # transform the offset to box coordinate
    Rs = get_R(boxes)  # (B, 3, 3)
    pc_box = Rs.transpose((0, 2, 1)) @ pc_box  # (B, 3, N)

    # check if points in box
    lwh_half = 0.5 * boxes[:, 3:6, np.newaxis] * resize_factor  # (B, 3, 1)
    xyz_mask = np.logical_and(pc_box >= -lwh_half, pc_box <= lwh_half)  # (B, 3, N)
    xyz_mask = xyz_mask.min(axis=1)  # xyz all in range, (B, N)
    in_box_mask = xyz_mask.max(axis=0)  # belong to at least one box (N, )

    # in the case of a point within two boxes, it is assigned to the closer one
    pc_box_norm = np.linalg.norm(pc_box, axis=1)  # (B, N)
    closest_box_inds = pc_box_norm.argmin(axis=0)  # (N, )

    return in_box_mask, closest_box_inds

def get_iou3d(boxes, query_boxes, need_bev=False):
    """Compute pair-wise IoU between two group of boxes.

    Modified from
    https://github.com/sshaoshuai/PointRCNN/blob/master/lib/utils/kitti_utils.py#L195

    Args:
        boxes (array[N, 7]): (x, y, z, l, w, h, theta)
        query_boxes (array[M, 7])
        need_bev (bool): True to also return IoU of BEV

    Returns:
        iou_3d (array[N, M])
        iou_bev (array[N, M])
    """
    N, M = boxes.shape[0], query_boxes.shape[0]
    corners_a = boxes_to_corners(boxes)  # (N, 3, 8)
    corners_b = boxes_to_corners(query_boxes)

    iou_3d = np.zeros((N, M), dtype=np.float32)
    iou_bev = np.zeros((N, M), dtype=np.float32)

    min_z_a = corners_a[:, 2, -1]
    max_z_a = corners_a[:, 2, 0]
    min_z_b = corners_b[:, 2, -1]
    max_z_b = corners_b[:, 2, 0]

    for i in range(N):
        for j in range(M):
            max_of_min = np.max([min_z_a[i], min_z_b[j]])
            min_of_max = np.min([max_z_a[i], max_z_b[j]])
            z_overlap = np.max([0, min_of_max - max_of_min])
            if z_overlap == 0:
                continue

            bottom_a = Polygon(corners_a[i, :2, 4:].T)
            bottom_b = Polygon(corners_b[j, :2, 4:].T)

            # check is valid,  A valid Polygon may not possess any overlapping
            # exterior or interior rings.
            if bottom_a.is_valid and bottom_b.is_valid:
                bottom_overlap = bottom_a.intersection(bottom_b).area
            else:
                bottom_overlap = 0.0

            overlap3d = bottom_overlap * z_overlap
            union3d = (
                bottom_a.area * (max_z_a[i] - min_z_a[i])
                + bottom_b.area * (max_z_b[j] - min_z_b[j])
                - overlap3d
            )
            iou_3d[i, j] = overlap3d / union3d
            iou_bev[i, j] = bottom_overlap / (
                bottom_a.area + bottom_b.area - bottom_overlap
            )

    if need_bev:
        return iou_3d, iou_bev

    return iou_3d

def get_iou3d_single(box_a, box_b, need_bev=False):
    """Compute IoU between two boxes.

    Args:
        box_a (array[7]): (x, y, z, l, w, h, theta)
        box_b (array[7]): (x, y, z, l, w, h, theta)
        need_bev (bool): True to also return IoU of BEV

    Returns:
        iou_3d (float)
    """
    return get_iou3d(box_a[np.newaxis], box_b[np.newaxis], need_bev)

def _boxes_jrdb_to_pointrcnn_convention(boxes):
    # NOTE don't really know why, but seems like iou3d expects KITTI box with
    # swapped x and z axis, otherwise cannot pass unittest test_iou3d() and test_nms3d()
    boxes_kitti = convert_boxes_jrdb_to_kitti_torch(boxes)
    tmp_r2 = boxes_kitti[:, 2].clone()
    boxes_kitti[:, 2] = boxes_kitti[:, 0]
    boxes_kitti[:, 0] = tmp_r2

    return boxes_kitti

def get_iou3d_gpu(boxes, query_boxes):
    """Compute pair-wise IoU between two group of boxes.

    Args:
        boxes (tensor[N, 7]): (x, y, z, l, w, h, theta)
        query_boxes (tensor[M, 7])

    Returns:
        iou_3d (tensor[N, M])
    """
    boxes = _boxes_jrdb_to_pointrcnn_convention(boxes)
    query_boxes = _boxes_jrdb_to_pointrcnn_convention(query_boxes)
    iou_3d = iou3d.boxes_iou3d_gpu(boxes, query_boxes)

    return iou_3d

def get_iou_bev_gpu(boxes, query_boxes):
    boxes = _boxes_jrdb_to_pointrcnn_convention(boxes)
    query_boxes = _boxes_jrdb_to_pointrcnn_convention(query_boxes)
    iou_bev = iou3d.boxes_iou_bev(
        iou3d.boxes3d_to_bev_torch(boxes), iou3d.boxes3d_to_bev_torch(query_boxes)
    )

    return iou_bev

def nms_3d(boxes, scores, iou_thresh=0.3):
    """IoU based NMS

    Args:
        boxes (array[B, 7]): (x, y, z, l, w, h, theta)
        scores (array[B]): cls score for each box
        iou_thresh (float):

    Returns:
        inds (array[N]): The inds of non-duplicate box, N <= B (see np.unique)
        inverse_map (array[B]): For each box, return the index of its parents
    """
    ious = get_iou3d(boxes, boxes)  # (B, B)
    sort_inds = np.argsort(scores)[::-1]

    inds = []
    inverse_map = -np.ones(ious.shape[0], dtype=np.int32)

    for i in sort_inds:
        if inverse_map[i] == -1:
            dup_mask = np.logical_and(ious[i] > iou_thresh, inverse_map == -1)
            inverse_map[dup_mask] = i
            inds.append(i)

    return np.array(inds, dtype=np.int32), inverse_map

def nms_3d_gpu(boxes, scores, iou_thresh=0.3):
    """IoU based NMS

    Args:
        boxes (tensor[B, 7]): (x, y, z, l, w, h, theta)
        scores (tensor[B]): cls score for each box
        iou_thresh (float):

    Returns:
        inds (tensor[N]): The inds of non-duplicate box, N <= B (see np.unique)
    """
    boxes_bev = iou3d.boxes3d_to_bev_torch(_boxes_jrdb_to_pointrcnn_convention(boxes))
    return iou3d.nms_gpu(boxes_bev, scores, iou_thresh)

def nms_3d_dist(boxes, scores, l_ave=0.9, w_ave=0.5, nms_thresh=0.4):
    """Distance based NMS

    Args:
        boxes (array[B, 7]): (x, y, z, l, w, h, theta)
        scores (array[B]): cls score for each box
        l_ave (float): average box length
        w_ave (float): average box width
        nms_thresh (float):

    Returns:
        inds (array[N]): The inds of non-duplicate box, N <= B (see np.unique)
        inverse_map (array[B]): For each box, return the index of its parents
    """
    x_diff = boxes[:, 0, np.newaxis] - boxes[np.newaxis, :, 0]
    y_diff = boxes[:, 1, np.newaxis] - boxes[np.newaxis, :, 1]
    d_diff = np.hypot(x_diff, y_diff)
    t_diff = np.arctan2(y_diff, x_diff)

    cos_t = np.abs(np.cos(boxes[:, 6, np.newaxis] - t_diff))  # this is not symmetric
    d_norm = w_ave + 0.5 * (cos_t + cos_t.T) * (l_ave - w_ave)

    d_ratio = d_diff / d_norm
    sort_inds = np.argsort(scores)[::-1]

    inds = []
    inverse_map = -np.ones(d_ratio.shape[0], dtype=np.int32)

    for i in sort_inds:
        if inverse_map[i] == -1:
            dup_mask = np.logical_and(d_ratio[i] < nms_thresh, inverse_map == -1)
            inverse_map[dup_mask] = i
            inds.append(i)

    return np.array(inds, dtype=np.int32), inverse_map

def nms_3d_dist_gpu(boxes, scores, l_ave=0.9, w_ave=0.5, nms_thresh=0.4):
    """IoU based NMS

    Args:
        boxes (tensor[B, 7]): (x, y, z, l, w, h, theta)
        scores (tensor[B]): cls score for each box
        l_ave (float): average box length
        w_ave (float): average box width
        nms_thresh (float):

    Returns:
        inds (tensor[N]): The inds of non-duplicate box, N <= B (see np.unique)
    """
    boxes_bev = iou3d.boxes3d_to_bev_torch(_boxes_jrdb_to_pointrcnn_convention(boxes))
    return iou3d.nms_dist_gpu(boxes_bev, scores, l_ave, w_ave, nms_thresh)

def _ravel_hash_vec(arr):
    assert arr.ndim == 2
    arr = arr - arr.min(0)
    arr = arr.astype(np.uint64, copy=False)
    arr_max = arr.max(0).astype(np.uint64) + 1

    keys = np.zeros(arr.shape[0], dtype=np.uint64)
    # Fortran style indexing
    for j in range(arr.shape[1] - 1):
        keys += arr[:, j]
        keys *= arr_max[j + 1]
    keys += arr[:, -1]
    return keys

def get_unique_rows(arr):
    keys = _ravel_hash_vec(arr)
    _, inds = np.unique(keys, return_index=True)
    return inds

class Box3d:
    def __init__(self, xyz, lwh, rot_z, box_id):
        """A 3D bounding box.

        Args:
            xyz (np.array(3,)): center xyz
            lwh (np.array(3,)): length, width, height. When rot_z is zero, length
                should be aligned with x axis
            rot_z (float): rotation around z axis
            box_id (int): label id used for tracking
        """
        self._xyz = xyz.reshape(-1)
        self._lwh = lwh.reshape(-1)
        self._rot_z = rot_z
        self._id = box_id

    def to_array(self):
        return np.concatenate([self._xyz, self._lwh, [self._rot_z]], axis=0)

    def get_xyz(self):
        return self._xyz

    def get_lwh(self):
        return self._lwh

    def get_id(self):
        return self._id

    def get_R(self):
        """See get_R()"""
        box = self.to_array()
        R = get_R(box[np.newaxis, :])

        return R[0]

    def to_corners(self, resize_factor=1.0, connect_inds=False):
        """See boxes_to_corners()"""
        box = self.to_array()

        if connect_inds:
            corners_xyz, connect_inds = boxes_to_corners(
                box[np.newaxis, :], resize_factor=resize_factor, connect_inds=True
            )
            return corners_xyz[0], connect_inds
        else:
            corners_xyz = boxes_to_corners(
                box[np.newaxis, :], resize_factor=resize_factor, connect_inds=False
            )
            return corners_xyz[0]

    def draw_bev(self, ax, c="red"):
        corner_xyz = self.to_corners()

        # side and back boarder
        xy = corner_xyz[:2, [1, 2, 3, 0]]
        ax.plot(xy[0], xy[1], c=c, linestyle="-")

        # front boarder
        xy = corner_xyz[:2, [0, 1]]
        ax.plot(xy[0], xy[1], c=c, linestyle="--")

    def draw_fpv(self, ax, dim, c="red"):
        """Plot first person view.

        Args:
            ax: Axes handle of matplotlib
            dim (int): 0 (x) for xz plot or 1 (y) for yz plot
            c (str, optional): color. Defaults to "red".
        """
        corner_xyz = self.to_corners()

        # top
        x = corner_xyz[dim, [0, 1, 2, 3, 0]]
        z = corner_xyz[2, [0, 1, 2, 3, 0]]
        ax.plot(x, z, c=c, linestyle="-")

        # bottom
        x = corner_xyz[dim, [4, 5, 6, 7, 4]]
        z = corner_xyz[2, [4, 5, 6, 7, 4]]
        ax.plot(x, z, c=c, linestyle="-")

        # vertical bar
        for i in range(4):
            x = corner_xyz[dim, [i, i + 4]]
            z = corner_xyz[2, [i, i + 4]]
            ax.plot(x, z, c=c, linestyle="-")

        # mark orientation
        x = corner_xyz[dim, [0, 5]]
        z = corner_xyz[2, [0, 5]]
        ax.plot(x, z, c=c, linestyle="--")
        x = corner_xyz[dim, [1, 4]]
        z = corner_xyz[2, [1, 4]]
        ax.plot(x, z, c=c, linestyle="--")

"""
Conversion API between JRDB and KITTI
"""

def convert_boxes_jrdb_to_kitti(jrdb_boxes):
    """Convert JRDB boxes to KITTI boxes

    x' = -y
    y' = -z + h / 2
    z' = x
    h' = h
    w' = w
    l' = l
    theta' = -theta

    Args:
        jrdb_boxes (array[B, 7]): x, y, z, l, w, h, theta

    Returns:
        kitti_boxes (array[B, 7]): x', y', z', h', w', l', theta'
    """
    # NOTE: Use test/test_iou3d() to check that this conversion is done correctly.
    kitti_boxes = np.stack(
        [
            -jrdb_boxes[:, 1],
            -jrdb_boxes[:, 2] + 0.5 * jrdb_boxes[:, 5],
            jrdb_boxes[:, 0],
            jrdb_boxes[:, 5],
            jrdb_boxes[:, 4],
            jrdb_boxes[:, 3],
            -jrdb_boxes[:, 6],
        ],
        axis=1,
    )

    return kitti_boxes

def convert_boxes_jrdb_to_kitti_torch(jrdb_boxes):
    kitti_boxes = torch.stack(
        [
            -jrdb_boxes[:, 1],
            -jrdb_boxes[:, 2] + 0.5 * jrdb_boxes[:, 5],
            jrdb_boxes[:, 0],
            jrdb_boxes[:, 5],
            jrdb_boxes[:, 4],
            jrdb_boxes[:, 3],
            -jrdb_boxes[:, 6],
        ],
        dim=1,
    )

    return kitti_boxes

def convert_boxes_kitti_to_jrdb(kitti_boxes):
    """Convert KITTI boxes to JRDB boxes

    x = z'
    y = -x'
    z = -y' + h' / 2
    l = l'
    w = w'
    h = h'
    theta = -theta'

    Args:
        kitti_boxes (array[B, 7]): x', y', z', h', w', l', theta'

    Returns:
        jrdb_boxes (array[B, 7]): x, y, z, l, w, h, theta
    """
    jrdb_boxes = np.stack(
        [
            kitti_boxes[:, 2],
            -kitti_boxes[:, 0],
            -kitti_boxes[:, 1] + 0.5 * kitti_boxes[:, 3],
            kitti_boxes[:, 5],
            kitti_boxes[:, 4],
            kitti_boxes[:, 3],
            -kitti_boxes[:, 6],
        ],
        axis=1,
    )

    return jrdb_boxes

def convert_boxes_kitti_to_jrdb_torch(kitti_boxes):
    jrdb_boxes = torch.stack(
        [
            kitti_boxes[:, 2],
            -kitti_boxes[:, 0],
            -kitti_boxes[:, 1] + 0.5 * kitti_boxes[:, 3],
            kitti_boxes[:, 5],
            kitti_boxes[:, 4],
            kitti_boxes[:, 3],
            -kitti_boxes[:, 6],
        ],
        dim=1,
    )

    return jrdb_boxes

"""
Dataset IO API
"""

def box_from_jrdb(jrdb_label, fast_mode=True):
    xyz_lwh_rot = np.array(
        [
            jrdb_label["box"]["cx"],
            jrdb_label["box"]["cy"],
            jrdb_label["box"]["cz"],
            jrdb_label["box"]["l"],
            jrdb_label["box"]["w"],
            jrdb_label["box"]["h"],
            jrdb_label["box"]["rot_z"],
        ],
        dtype=np.float32,
    )
    box_id = int(jrdb_label["label_id"].split(":")[-1])

    if fast_mode:
        return xyz_lwh_rot, box_id
    else:
        return Box3d(xyz_lwh_rot[:3], xyz_lwh_rot[3:6], xyz_lwh_rot[6], box_id)

def box_from_nuscenes(ann, fast_mode=True):
    xyz = ann["translation"]
    wlh = ann["size"]
    qt = Quaternion(ann["rotation"])
    # Rotation is only along vertical axis, but the rotation axis direction is
    # arbitary (plus or minus)
    r = -qt.axis[2] * qt.radians  # the minus sign is confirmed by plotting
    if abs(qt.axis[2]) != 1.0 and np.abs(qt.axis).sum() != 0.0:
        print(ann)
    box_id = ann["instance_token"]

    xyz_lwh_rot = np.array([*xyz, wlh[1], wlh[0], wlh[2], r], dtype=np.float32)

    if fast_mode:
        return xyz_lwh_rot, box_id
    else:
        return Box3d(xyz_lwh_rot[:3], xyz_lwh_rot[3:6], xyz_lwh_rot[6], box_id)

def box_to_nuscenes(box, score, sample_token, detection_name):
    xyz = box[:3].tolist()
    wlh = box[[4, 3, 5]].tolist()
    # the minus sign should be consistent with box_from_nuscenes()
    qt = Quaternion(axis=[0, 0, 1], angle=-box[6], dtype=float)

    sr_dict = {
        "sample_token": sample_token,
        "translation": xyz,
        "size": wlh,
        "rotation": list(qt),
        "velocity": [0.0, 0.0],
        "detection_name": detection_name,
        "detection_score": float(score),
        "attribute_name": "",
    }

    return sr_dict

def boxes_to_string(boxes, scores, jrdb_format=True):
    """Obtain a KITTI format string for storing detection result

    Args:
        boxes (array[B, 7])
        scores (array[B])
        jrdb_format (bool): True if the input boxes is in JRDB convention, False
            if in KITTI convention

    Returns:
        s (str)
    """
    if jrdb_format:
        if isinstance(boxes, np.ndarray):
            boxes = convert_boxes_jrdb_to_kitti(boxes)
        else:
            boxes = convert_boxes_jrdb_to_kitti_torch(boxes)

    s = ""
    for box, score in zip(boxes, scores):
        s += (
            f"Pedestrian 0 0 -1 0 -1 -1 -1 -1 {box[3]} {box[4]} {box[5]} "
            f"{box[0]} {box[1]} {box[2]} {box[6]} {score}\n"
        )
    s = s.strip("\n")

    return s

def string_to_boxes(s, jrdb_format=True, get_num_points=False):
    """Convert a KITTI format string to boxes and detection scores

    Args:
        s (str)
        jrdb_format (bool): True if the output boxes should be in JRDB convention,
            False if in KITTI convention
        get_num_points (bool): True to also return number of points. Useful
            for loading annotations

    Returns:
        boxes (array[B, 7])
        scores (array[B])
    """
    boxes = []
    scores = []

    lines = s.split("\n")
    for line in lines:
        if len(line) == 0:
            continue
        v_list = [float(v) for v in line.split()[-8:]]
        scores.append(v_list[-1])
        boxes.append([v_list[i] for i in [3, 4, 5, 0, 1, 2, 6]])

    boxes = np.array(boxes, dtype=np.float32)
    scores = np.array(scores, dtype=np.float32)

    if jrdb_format and len(boxes) > 0:
        boxes = convert_boxes_kitti_to_jrdb(boxes)

    if get_num_points:
        num_points = np.array([int(line.split()[3]) for line in lines if len(line) > 0])
        return boxes, scores, num_points
    else:
        return boxes, scores


# Credit 2017
# Brian D. Bue (brian.d.bue@jpl.nasa.gov)
# David R. Thompson (david.r.thompson@jpl.nasa.gov)
#
# License: Apache 2.0 (http://www.apache.org/licenses/)
#

from __future__ import division, print_function

import warnings
#warnings.filterwarnings("ignore", category=DeprecationWarning)
warn = warnings.warn

from spectral.io.envi import open as envi_open_file
from spectral.io.envi import create_image as envi_create_image

import time
from os import makedirs
from os.path import join as pathjoin, exists as pathexists
from numpy import pi, array, arange, diag, mean, median, zeros, ones
from numpy import all, any, isfinite, argmin, log, loadtxt
from numpy import uint8,int16,float32,float64
from numpy import matrix,unique,isfinite,inf,where
from numpy import logical_and
from scipy.linalg import LinAlgError

from sklearn.cluster import MiniBatchKMeans
from sklearn.decomposition import PCA

import pandas as pd

gettime = time.time

modelname='looshrinkage' # 'empirical'

ppmscaling = 100000.0

def randperm(*args):
    from numpy.random import permutation
    n = args[0]
    k = n if len(args) < 2 else args[1]
    return permutation(n)[:k]

def np2envitype(np_dtype):
    from numpy import dtype
    from spectral.io.envi import dtype_to_envi
    _dtype = dtype(np_dtype).char
    return dtype_to_envi[_dtype]

def cov(A,**kwargs):
    """
    cov(A,**kwargs)

    Summary: computes covariance that matches matlab covariance function (ddof=1)

    Arguments:
    - A: n x m array of n samples with m features per sample

    Keyword Arguments:
    - same as numpy.cov

    Output:
    m x m covariance matrix
    """

    from numpy import cov as _cov
    kwargs.setdefault('ddof',1)
    return _cov(A.T,**kwargs)

def inv(A,**kwargs):
    from scipy.linalg import inv as _inv
    kwargs.setdefault('overwrite_a',False)
    kwargs.setdefault('check_finite',False)
    return _inv(A,**kwargs)

def eig(A,**kwargs):
    from scipy.linalg import eig as _eig
    kwargs.setdefault('overwrite_a',False)
    kwargs.setdefault('check_finite',False)
    kwargs.setdefault('left',False)
    kwargs.setdefault('right',True)
    return _eig(A,**kwargs)

def det(A,**kwargs):
    from scipy.linalg import det as _det
    kwargs.setdefault('overwrite_a',False)
    kwargs.setdefault('check_finite',False)
    return _det(A,**kwargs)

def looshrinkage(I_zm,alphas,nll,n,I_reg=[]):
    # loocv shrinkage estimation via Theiler et al.
    stability_scaling=100.0
    nchan = I_zm.shape[1]

    X = I_zm*stability_scaling
    S = cov(X)
    T = diag(diag(S)) if len(I_reg)==0 else cov(I_reg*stability_scaling)

    nchanlog2pi = nchan*log(2.0*pi)
    nll[:] = inf

    # Closed form for leave one out cross validation error
    for i,alpha in enumerate(alphas):
        try:
            # See Theiler, "The Incredible Shrinking Covariance Estimator",
            # Proc. SPIE, 2012. eqn. 29
            beta = (1.0-alpha) / (n-1.0)
            G_alpha = n * (beta*S) + (alpha*T)
            G_det = det(G_alpha)
            if G_det==0:
                continue
            r_k  = (X.dot(inv(G_alpha)) * X).sum(axis=1)
            q = 1.0 - beta * r_k
            nll[i] = 0.5*(nchanlog2pi+log(G_det))+1.0/(2.0*n) * \
                     (log(q)+(r_k/q)).sum()
        except LinAlgError:
            warn('looshrinkage encountered a LinAlgError')

    mindex = argmin(nll)

    if nll[mindex]!=inf:
        alpha = alphas[mindex]
    else:
        mindex = -1
        alpha = 0.0

    # Final nonregularized covariance and shrinkage target
    S = cov(I_zm)
    T = diag(diag(S)) if len(I_reg)==0 else cov(I_reg)

    # Final covariance
    C = (1.0 - alpha) * S + alpha * T

    return C,mindex

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description="Robust MF")
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='verbose output')
    parser.add_argument('-k', '--kmeans', type=int, default=1,
                        help='number of columnwise modes (k-means clusters)')
    parser.add_argument('--pcadim', type=int, default=6,
                        help='number of PCA dims (for k-means clusters>1)')
    parser.add_argument('-r', '--reject', action='store_true',
                        help='enable multimodal covariance outlier rejection')
    parser.add_argument('-f', '--full', action='store_true',
                        help='regularize multimodal estimates with the full column covariariance')
    parser.add_argument('--rgb_bands', default='60,42,24',
                        help='comma-separated list of RGB channels')
    parser.add_argument('-m', '--metadata', action='store_true',
                        help='save metadata image')
    parser.add_argument('-R', '--reflectance', action='store_true',
                        help='reflectance signature')
    parser.add_argument('-M', '--model', type=str, default='looshrinkage',
                        help='model name (looshrinkage (default)|empirical)')
    parser.add_argument('input', type=str,  metavar='INPUT',
                        help='path to input image')
    parser.add_argument('library', type=str,  metavar='LIBRARY',
                        help='path to target library file')
    parser.add_argument('output', type=str,  metavar='OUTPUT',
                        help='path for output image (mf ch4 ppm)')

    args = parser.parse_args()
    bgmodes = args.kmeans
    reject = args.reject
    regfull = args.full
    reflectance = args.reflectance
    savebgmeta = args.metadata
    modelname = args.model
    pcadim = args.pcadim

    infile = args.input
    infilehdr = infile+'.hdr'
    outfile = args.output
    outfilehdr = outfile+'.hdr'

    # csv output file for colunnwise mean/stddev
    colcsv = os.path.splitext(infile)[0]+'_column_stats.csv'

    # define active channels (AVIRIS-NG clipped radiances) wrt target gas + measurement units
    if reflectance and 'ch4' in args.library:
      active = [5,420]
    elif 'ch4' in args.library:
      active = [351,422]
    elif 'co2' in args.library:
      active = [309,391]
    else:
      print('could not set active range')
      sys.exit(0)

    # columnwise spectral averaging function
    colavgfn = mean

    # want bg modes to have at least as many samples as 120% x (# features)
    bgminsamp = int((active[1]-active[0])*1.2)
    bgmodel = 'unimodal' if bgmodes==1 else 'multimodal'

    print('started processing input file: "%s"'%str(infile))

    img = envi_open_file(infilehdr,image=infile)
    img_mm = img.open_memmap(interleave='source',writeable=False)
    nrows,nbands,ncols = img_mm.shape

    outmeta = img.metadata
    outmeta['lines'] = nrows
    outmeta['data type'] = np2envitype(float64)

    rgb_bands = args.rgb_bands
    rgb_bands = [] if rgb_bands=='[]' else map(int,rgb_bands.split(','))
    if len(rgb_bands)==3:
        outmeta['bands'] = 4
        outmeta['band names'] = ['Red Radiance (uW/nm/sr/cm2)',
                                 'Green Radiance (uW/nm/sr/cm2)',
                                 'Blue Radiance (uW/nm/sr/cm2)',
                                 'CH4 Absorption (ppm x m)']
    elif len(rgb_bands)==0:
        outmeta['bands'] = 1
        outmeta['band names'] = ['CH4 Absorption (ppm x m)']
    else:
        raise Exception('invalid value of rgb_bands argument: %s'%args.rgb_bands)

    outmeta['interleave'] = 'bip'
    for kwarg in ['smoothing factors','wavelength','wavelength units','fwhm']:
        outmeta.pop(kwarg,None)

    nodata = float(outmeta.get('data ignore value',-9999))
    if nodata > 0:
        raise Exception('nodata value=%f > 0, values will not be masked'%nodata)

    # load the gas spectrum
    libdata = float64(loadtxt(args.library))
    abscf=libdata[active[0]-1:active[1],2]

    # alphas for leave-one-out cross validation shrinkage
    if modelname == 'looshrinkage':
        astep,aminexp,amaxexp = 0.05,-10.0,0.0
        alphas=(10.0 ** arange(aminexp,amaxexp+astep,astep))
        nll=zeros(len(alphas))

    modelparms  = 'modelname={modelname}, bgmodel={bgmodel}'
    if bgmodes > 1:
        modelparms += ', bgmodes={bgmodes}, pcadim={pcadim}, reject={reject}'
        if modelname == 'looshrinkage':
            modelparms += ', regfull={regfull}'

    if modelname == 'looshrinkage':
        modelparms += ', aminexp={aminexp}, amaxexp={amaxexp}, astep={astep}'

    modelparms += ', reflectance={reflectance}, active_bands={active}'

    outdict = locals()
    outdict.update(globals())
    outmeta['model parameters'] = '{ %s }'%(modelparms.format(**outdict))

    outimg = envi_create_image(outfilehdr,outmeta,force=True,ext='')
    outimg_mm = outimg.open_memmap(interleave='source',writable=True)
    assert((outimg_mm.shape[0]==nrows) & (outimg_mm.shape[1]==ncols))

    # outimage mf values = nodata by default
    outimg_mm[:,:,-1] = nodata

    if savebgmeta:
        # output image of bg cluster membership labels per column
        bgfile=outfile+'_bgmeta'
        bgfilehdr=bgfile+'.hdr'
        bgmeta = outmeta
        bgmeta['bands'] = 2
        bgmeta['data type'] = np2envitype(int16)
        bgmeta['num alphas'] = len(alphas)
        bgmeta['alphas'] = '{%s}'%(str(alphas)[1:-1])
        bgmeta['band names'] = '{cluster_id, alpha_index}'
        bgimg = envi_create_image(bgfilehdr,bgmeta,force=True,ext='')
        bgimg_mm = bgimg.open_memmap(interleave='source',writable=True)

    # exclude nonfinite + negative spectra in covariance estimates
    useidx = lambda Icol: where(((~(Icol<0)) & isfinite(Icol)).all(axis=1))[0]

    # if len(rgb_bands)==3:
    #     print('copying rgb bands to outfile')
    #     stime = gettime()
    #     for oi,bi in enumerate(rgb_bands):
    #         outimg_mm[:,:,oi] = img_mm[:,bi,:]
    #     print('done (elapsed time=%ds)'%(gettime()-stime))

    print('starting columnwise processing (%d columns)'%ncols)
    stime = gettime()
    colavg = np.ones(ncols)*nodata
    colstd = np.ones(ncols)*nodata
    colnum = np.ones(ncols)*nodata

    for col in arange(ncols):
        Icol_full=img_mm[:,active[0]-1:active[1],col]
        use = useidx(Icol_full)

        Icol = float64(Icol_full[use,:].copy())
        nuse = Icol.shape[0]
        if nuse == 0:
            continue

        if bgmodes > 1:
            # PCA projection down to a smaller number of dimensions
            # then apply K-means to separate spatially into clusters
            Icol_zm = Icol-colavgfn(Icol,axis=0)
            evals,evecs = eig(cov(Icol_zm))
            Icol_pca = Icol_zm.dot(evecs[:,:pcadim])
            cmodel = MiniBatchKMeans(n_clusters=bgmodes)
            bglabels = cmodel.fit(Icol_pca).labels_
            bgulab = unique(bglabels)

            bgcounts = []
            bgulabn = zeros(len(bgulab))
            for i,l in enumerate(bgulab):
                lmask = bglabels==l
                bgulabn[i] = lmask.sum()
                if reject and bgulabn[i] < bgminsamp:
                    print('flagged outlier cluster %d (%d samples)'%(l,bgulabn[i]))
                    bglabels[lmask] = -l
                    bgulab[i] = -l
                bgcounts.append("%d: %d"%(l,bgulabn[i]))
                if savebgmeta:
                    bgimg_mm[use[lmask],col,0] = bgulab[i]

            print('bg cluster counts:',', '.join(bgcounts))
            if (bgulab<0).all():
                warn('all clusters rejected, proceeding without rejection (beware!)')
                bglabels,bgulab = abs(bglabels),abs(bgulab)

        else: # bgmodes==1
            bglabels = ones(nuse)
            bgulab = array([1])

        # operate independently on each columnwise partition
        for ki in bgulab:
            # if bglabel<0 (=rejected), estimate using all (nonrejected) modes
            kmask = bglabels==ki if ki >= 0 else bglabels>=0

            # need to recompute mu and associated vars wrt this cluster
            Icol_ki = (Icol if bgmodes == 1 else Icol[kmask,:]).copy()

            Icol_sub = Icol_ki.copy()
            mu = colavgfn(Icol_sub,axis=0)

            # reinit model/modelfit here for each column/cluster instance
            if modelname == 'empirical':
                modelfit = lambda I_zm: cov(I_zm)
            elif modelname == 'looshrinkage':
                # optionally use the full zero mean column as a regularizer
                Icol_reg = Icol-mu if (regfull and bgmodes>1) else []
                modelfit = lambda I_zm: looshrinkage(I_zm,alphas,nll,
                                                     nuse,I_reg=Icol_reg)

            try:
                Icol_sub = Icol_sub-mu # = subsampled column mode
                Icol_model = modelfit(Icol_sub)
                if modelname=='looshrinkage':
                    C,alphaidx = Icol_model
                    Cinv=inv(C)
                    if savebgmeta:
                        bgimg_mm[use[kmask],col,1] = alphaidx
                elif modelname=='empirical':
                    Cinv = inv(Icol_model)
                else:
                    Cinv = Icol_model

            except LinAlgError:
                print('singular matrix. skipping this column mode.')
                outimg_mm[use[kmask],col,-1] = 0
                continue

            # Classical matched filter
            Icol_ki = Icol_ki-mu # = fully-sampled column mode
            target = abscf.copy()
            target = target-mu if reflectance else target*mu
            normalizer = target.dot(Cinv).dot(target.T)
            mf = (Icol_ki.dot(Cinv).dot(target.T)) / normalizer

            if reflectance:
                outimg_mm[use[kmask],col,-1] = mf
            else:
                outimg_mm[use[kmask],col,-1] = mf*ppmscaling

        colpix = outimg_mm[use[bglabels>=0],col,-1]
        colnum[col] = nuse
        colavg[col] = np.mean(colpix)
        colstd[col] = np.std(colpix)
        print('Column %i mean: %e, std: %e'%(col,colavg[col],colstd[col]))

        # copy rgb bands to outfile
        outimg_mm[:,col,0] = img_mm[:,rgb_bands[0],col]
        outimg_mm[:,col,1] = img_mm[:,rgb_bands[1],col]
        outimg_mm[:,col,2] = img_mm[:,rgb_bands[2],col]

    # write column stats
    print('Saving column stats to',colcsv)
    coldf = pd.DataFrame(np.r_[colnum,colavg,colstd],
                         index=['npix','avg','std'])
    coldf.to_csv(colcsv)

    print('done (elapsed time=%ds)'%(gettime()-stime))


import logging
import sys
from typing import Dict

from ml.rl.evaluation.evaluator import Evaluator
from ml.rl.preprocessing.normalization import (
    construct_action_scale_tensor,
    get_num_output_features,
)
from ml.rl.preprocessing.preprocessor import Preprocessor
from ml.rl.preprocessing.sparse_to_dense import PandasSparseToDenseProcessor
from ml.rl.readers.json_dataset_reader import JSONDatasetReader
from ml.rl.tensorboardX import summary_writer_context
from ml.rl.thrift.core.ttypes import (
    ContinuousActionModelParameters,
    DDPGModelParameters,
    DDPGNetworkParameters,
    DDPGTrainingParameters,
    NormalizationParameters,
    RLParameters,
)
from ml.rl.training.ddpg_trainer import ActorNetModel, CriticNetModel, DDPGTrainer
from ml.rl.training.rl_exporter import ActorExporter
from ml.rl.workflow.base_workflow import BaseWorkflow
from ml.rl.workflow.helpers import (
    export_trainer_and_predictor,
    minibatch_size_multiplier,
    parse_args,
    update_model_for_warm_start,
)
from ml.rl.workflow.preprocess_handler import (
    ContinuousPreprocessHandler,
    PreprocessHandler,
)
from tensorboardX import SummaryWriter

logger = logging.getLogger(__name__)

class ContinuousWorkflow(BaseWorkflow):
    def __init__(
        self,
        model_params: ContinuousActionModelParameters,
        preprocess_handler: PreprocessHandler,
        state_normalization: Dict[int, NormalizationParameters],
        action_normalization: Dict[int, NormalizationParameters],
        use_gpu: bool,
        use_all_avail_gpus: bool,
    ):
        logger.info("Running continuous workflow with params:")
        logger.info(model_params)
        model_params = model_params

        min_action_range_tensor_serving, max_action_range_tensor_serving = construct_action_scale_tensor(
            action_normalization, model_params.action_rescale_map
        )

        state_dim = get_num_output_features(state_normalization)
        action_dim = get_num_output_features(action_normalization)

        # Build Actor Network
        actor_network = ActorNetModel(
            layers=(
                [state_dim] + model_params.actor_training.layers[1:-1] + [action_dim]
            ),
            activations=model_params.actor_training.activations,
            fl_init=model_params.shared_training.final_layer_init,
            state_dim=state_dim,
            action_dim=action_dim,
            use_gpu=use_gpu,
            use_all_avail_gpus=use_all_avail_gpus,
        )

        # Build Critic Network
        critic_network = CriticNetModel(
            # Ensure dims match input state and scalar output
            layers=[state_dim] + model_params.critic_training.layers[1:-1] + [1],
            activations=model_params.critic_training.activations,
            fl_init=model_params.shared_training.final_layer_init,
            state_dim=state_dim,
            action_dim=action_dim,
            use_gpu=use_gpu,
            use_all_avail_gpus=use_all_avail_gpus,
        )

        trainer = DDPGTrainer(
            actor_network,
            critic_network,
            model_params,
            state_normalization,
            action_normalization,
            min_action_range_tensor_serving,
            max_action_range_tensor_serving,
            use_gpu=use_gpu,
            use_all_avail_gpus=use_all_avail_gpus,
        )
        trainer = update_model_for_warm_start(trainer)
        assert type(trainer) == DDPGTrainer, "Warm started wrong model type: " + str(
            type(trainer)
        )

        evaluator = Evaluator(
            None,
            model_params.rl.gamma,
            trainer,
            metrics_to_score=trainer.metrics_to_score,
        )

        super().__init__(
            preprocess_handler,
            trainer,
            evaluator,
            model_params.shared_training.minibatch_size,
        )

def _get_actor_exporter(trainer, state_normalization, action_normalization):
    return ActorExporter.from_state_action_normalization(
        trainer.actor,
        state_normalization=state_normalization,
        action_normalization=action_normalization,
    )

def main(params):
    # Set minibatch size based on # of devices being used to train
    params["shared_training"]["minibatch_size"] *= minibatch_size_multiplier(
        params["use_gpu"], params["use_all_avail_gpus"]
    )

    rl_parameters = RLParameters(**params["rl"])
    training_parameters = DDPGTrainingParameters(**params["shared_training"])
    actor_parameters = DDPGNetworkParameters(**params["actor_training"])
    critic_parameters = DDPGNetworkParameters(**params["critic_training"])

    model_params = DDPGModelParameters(
        rl=rl_parameters,
        shared_training=training_parameters,
        actor_training=actor_parameters,
        critic_training=critic_parameters,
    )

    state_normalization = BaseWorkflow.read_norm_file(params["state_norm_data_path"])
    action_normalization = BaseWorkflow.read_norm_file(params["action_norm_data_path"])

    writer = SummaryWriter(log_dir=params["model_output_path"])
    logger.info("TensorBoard logging location is: {}".format(writer.log_dir))

    preprocess_handler = ContinuousPreprocessHandler(
        Preprocessor(state_normalization, False),
        Preprocessor(action_normalization, False),
        PandasSparseToDenseProcessor(),
    )

    workflow = ContinuousWorkflow(
        model_params,
        preprocess_handler,
        state_normalization,
        action_normalization,
        params["use_gpu"],
        params["use_all_avail_gpus"],
    )

    train_dataset = JSONDatasetReader(
        params["training_data_path"], batch_size=training_parameters.minibatch_size
    )
    eval_dataset = JSONDatasetReader(params["eval_data_path"], batch_size=16)

    with summary_writer_context(writer):
        workflow.train_network(train_dataset, eval_dataset, int(params["epochs"]))
    return export_trainer_and_predictor(
        workflow.trainer,
        params["model_output_path"],
        exporter=_get_actor_exporter(
            trainer=workflow.trainer,
            state_normalization=state_normalization,
            action_normalization=action_normalization,
        ),
    )  # noqa

if __name__ == "__main__":
    logging.basicConfig(stream=sys.stdout, level=logging.INFO)
    params = parse_args(sys.argv)
    main(params)


import numpy
import copy
import matplotlib
from matplotlib import pyplot as plt
import torch
import torch.nn as nn

from utils import set_up_dir
import os
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy

def get_prediction(model, inputs, device):

    with torch.no_grad():
        inputs = inputs.to(device)
        outputs = model(inputs)
        probs = nn.Softmax(dim=-1)(outputs)
        vals, preds = torch.max(outputs, 1)

    return int(preds.squeeze()), outputs.detach().cpu().numpy().squeeze(), probs.detach().cpu().numpy().squeeze()

def heatmap(R, fax = None):

    b = 10*((numpy.abs(R)**3.0).mean()**(1.0/3))
    from matplotlib.colors import ListedColormap
    my_cmap = plt.cm.seismic(numpy.arange(plt.cm.seismic.N))
    my_cmap[:,0:3] *= 0.85
    my_cmap = ListedColormap(my_cmap)
    f,ax = fax

    plt.subplots_adjust(left=0,right=1,bottom=0,top=1)
    ax.axis('off')
    h=ax.imshow(R,cmap=my_cmap,vmin=-b,vmax=b,interpolation='nearest')
    plt.colorbar(h)

def newlayer(layer,g):

    layer = copy.deepcopy(layer)

    try: layer.weight = nn.Parameter(g(layer.weight))
    except AttributeError: pass

    try: layer.bias   = nn.Parameter(g(layer.bias))
    except AttributeError: pass

    return layer

def toconv(layers):
    # convert VGG classifier's dense layers to convolutional layers

    newlayers = []

    for i,layer in enumerate(layers):
        if isinstance(layer,nn.Linear):
            newlayer = None
            if i == 0:
                m,n = 512,layer.weight.shape[0]
                newlayer = nn.Conv2d(m,n,7)
                newlayer.weight = nn.Parameter(layer.weight.reshape(n,m,7,7))
            else:
                m,n = layer.weight.shape[1],layer.weight.shape[0]
                newlayer = nn.Conv2d(m,n,1)
                newlayer.weight = nn.Parameter(layer.weight.reshape(n,m,1,1))

            newlayer.bias = nn.Parameter(layer.bias)
            newlayers += [newlayer]
        else:
            newlayers += [layer]

    return newlayers

def explain(model, X, cl, device='cpu', n_classes=2):
    # compute LRP explanations
    # https://git.tu-berlin.de/gmontavon/lrp-tutorial

    layers = list(model._modules['features']) +  toconv([model._modules['avgpool']]) +toconv(list(model._modules['classifier']))
    L = len(layers)

    A = [X]+[None]*L
    for l in range(L):
        A[l+1] = layers[l].forward(A[l])

    model_outputs = model(X)

    T = torch.FloatTensor((1.0*(numpy.arange(n_classes)==cl).reshape([1,n_classes,1,1]))).to(device)

    R = [None]*L + [(A[-1]*T).data]

    for l in range(1,L)[::-1]:

        A[l] = (A[l].data).requires_grad_(True)

        if isinstance(layers[l],torch.nn.MaxPool2d): layers[l] = torch.nn.AvgPool2d(2)

        if isinstance(layers[l],torch.nn.Conv2d) or isinstance(layers[l],torch.nn.AvgPool2d) or isinstance(layers[l],torch.nn.AdaptiveAvgPool2d) :

            if l <= 10:       rho = lambda p: p + 0.5*p.clamp(min=0);  incr = lambda z: z+1e-9
            if 11 <= l <= 17: rho = lambda p: p + 0.25*p.clamp(min=0); incr = lambda z: z+1e-9
            if 18 <= l <= 24: rho = lambda p: p + 0.1*p.clamp(min=0);  incr = lambda z: z+1e-9
            if l > 24:       rho = lambda p: p;                        incr = lambda z: z+1e-9

            z = incr(newlayer(layers[l],rho).forward(A[l]))        # step 1
            s = (R[l+1]/z).data                                    # step 2
            (z*s).sum().backward(); c = A[l].grad                  # step 3
            R[l] = (A[l]*c).data                                   # step 4

        else:
            R[l] = R[l+1]

    mean = torch.Tensor([0.5, 0.5, 0.5]).reshape(1,-1,1,1).to(device)
    std  = torch.Tensor([0.5, 0.5, 0.5]).reshape(1,-1,1,1).to(device)

    A[0] = (A[0].data).requires_grad_(True)

    lb = (A[0].data*0+(0-mean)/std).requires_grad_(True)
    hb = (A[0].data*0+(1-mean)/std).requires_grad_(True)

    z = layers[0].forward(A[0]) + 1e-9                                     # step 1 (a)
    z -= newlayer(layers[0],lambda p: p.clamp(min=0)).forward(lb)          # step 1 (b)
    z -= newlayer(layers[0],lambda p: p.clamp(max=0)).forward(hb)          # step 1 (c)
    s = (R[1]/z).data                                                      # step 2
    (z*s).sum().backward(); c,cp,cm = A[0].grad,lb.grad,hb.grad            # step 3
    R[0] = (A[0]*c+lb*cp+hb*cm).data

    return R

def plot_heatmap(R, x, fax, l, title_map, cl, b):

    f,ax = fax
    from matplotlib.colors import ListedColormap
    from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable

    my_cmap = plt.cm.seismic(numpy.arange(plt.cm.seismic.N))
    my_cmap[:,0:3] *= 0.85
    my_cmap = ListedColormap(my_cmap)

    # plt.figure(figsize=(sx,sy))
    plt.subplots_adjust(left=0,right=1,bottom=0,top=1)
    ax.axis('off')
    ax.imshow(x, cmap='Greys')
    if l==0:
        ax.set_title('explained : {}'.format(title_map[cl]))
    else:
        ax.set_title('{}'.format(title_map[cl]))

    h=ax.imshow(R,cmap=my_cmap,vmin=-b,vmax=b,interpolation='nearest', alpha=0.9)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.05)

    if l!=len(title_map)-1:
        cax.set_axis_off()
    else:
        plt.colorbar(h, cax=cax)


import json
from django.conf import settings
from django.urls import reverse

from .models import XMPPAccount

def get_conversejs_settings(settings_dict=None):
    converse_settings = {
        'bosh_service_url': settings.XMPP_BOSH_SERVICE_URL,
        'domain_placeholder': settings.XMPP_DOMAIN,
        'debug': settings.DEBUG,
        'authentication': getattr(settings, 'XMPP_CONVERSEJS_AUTH', 'login'),
        'keepalive': True,
        'auto_login': True,
    }

    try:
        converse_settings.update(settings.XMPP_CONVERSEJS_SETTINGS)
    except AttributeError:
        pass

    if settings:
        converse_settings.update(settings_dict)

    if converse_settings.get('xhr_user_search', None):
        converse_settings['xhr_user_search_url'] = reverse('xmpp_xhr_user_search')

    return converse_settings

def get_conversejs_context(context):
    context['xmpp_enabled'] = settings.XMPP_ENABLED

    if not settings.XMPP_ENABLED:
        # XMPP disabled entirely
        return context

    request = context.get('request')
    if not request.user.is_active:
        # User is not logged in
        return context

    # Setup authentication
    auth_type = getattr(settings, 'XMPP_CONVERSEJS_AUTH', None)
    settings_auth = {}
    if auth_type in ['login', 'prebind']:
        # We need to ensure XMPPAccount for both login and prebind auth
        xmpp_account = XMPPAccount.get_or_create(request.user)
        settings_auth['jid'] = xmpp_account.jid

    if auth_type == 'login':
        settings_auth['password'] = xmpp_account.password
    elif auth_type == 'preauth':
        raise NotImplementedError('preauth is not implemented yet')

    context['conversejs_settings'] = json.dumps(get_conversejs_settings(settings_auth))

    return context


"""server URL Configuration

The `urlpatterns` list routes URLs to views. For more information please see:
    https://docs.djangoproject.com/en/3.1/topics/http/urls/
Examples:
Function views
    1. Add an import:  from my_app import views
    2. Add a URL to urlpatterns:  path('', views.home, name='home')
Class-based views
    1. Add an import:  from other_app.views import Home
    2. Add a URL to urlpatterns:  path('', Home.as_view(), name='home')
Including another URLconf
    1. Import the include() function: from django.urls import include, path
    2. Add a URL to urlpatterns:  path('blog/', include('blog.urls'))
"""
from django.contrib import admin
from django.urls import path

from newfirst import views

urlpatterns = [
    path('item_list', views.item_list, name='item_list'),  # 项目列表
    path('design_criteria_list', views.design_criteria_list, name='design_criteria_list'),  # 设计准则列表
    path('analysis_rule_list', views.analysis_rule_list, name='analysis_rule_list'),  # 分析规则列表
    path('add_analysis_rule', views.add_analysis_rule, name='add_analysis_rule'),  # 添加分析规则
    path('add_analysis_rule_from_item', views.add_analysis_rule_from_item, name='add_analysis_rule_from_item'),
    path('edit_analysis_rule', views.edit_analysis_rule, name='edit_analysis_rule'),  # 编辑分析规则
    path('delete_analysis_rule', views.delete_analysis_rule, name='delete_analysis_rule'),  # 删除分析规则
    path('add_design_criteria', views.add_design_criteria, name='add_design_criteria'),  # 添加设计准则
    path('add_design_criteria_from_item', views.add_design_criteria_from_item, name='add_design_criteria_from_item'),
    path('edit_design_criteria', views.edit_design_criteria, name='edit_design_criteria'),  # 编辑设计准则
    path('delete_design_criteria', views.delete_design_criteria, name='delete_design_criteria'),  # 删除设计准则
    path('add_item', views.add_item, name='add_item'),  # 添加项目
    path('edit_item', views.edit_item, name='edit_item'),  # 编辑项目
    path('delete_item', views.delete_item, name='delete_item'),  # 删除项目
    path('scenes_list', views.scenes_list, name='scenes_list'),  # 场景列表
    path('add_scenes', views.add_scenes, name='add_scenes'),  # 添加场景
    path('edit_scenes', views.edit_scenes, name='edit_scenes'),  # 编辑场景
    path('delete_scenes', views.delete_scenes, name='delete_scenes'),  # 删除场景
    path('add_rule', views.add_rule, name='add_rule'),  # 添加规则
    path('rules_list', views.rules_list, name='rules_list'),  # 规则集列表
    path('delete_rule', views.delete_rule, name='delete_rule'),  # 删除规则
    path('add_case_list', views.add_case_list, name='add_case_list'),  # 添加实例列表
    path('add_case', views.add_case, name='add_case'),  # 添加实例
    path('case_list', views.case_list, name='case_list'),  # 添加实例列表
    path('delete_case', views.delete_case, name='delete_case'),  # 删除实例
    path('verify_case', views.verify_case, name='verify_case'),  # 验证实例
    path('verify_case_test', views.verify_case_test, name='verify_case_test'),  # 预检验
    path('reset_case', views.reset_case, name='reset_case'),  # 重置实例
    path('fmea_list', views.fmea_list, name='fmea_list'),  # fmea列表
    path('edit_fmea', views.edit_fmea, name='edit_fmea'),  # 编辑fmea
    path('demand_list', views.demand_list, name='demand_list'),  # 需求表
    path('edit_demand', views.edit_demand, name='edit_demand'),  # 编辑需求表
    path('add_design', views.add_design, name='add_design'),  # 添加准则库
    path('designs_list', views.designs_list, name='designs_list'),  # 规则集列表
    path('delete_design', views.delete_design, name='delete_design'),  # 删除规则
    path('check_list', views.check_list, name='check_list'),  # 设计核查
    path('edit_check', views.edit_check, name='edit_check'),  # 编辑设计核查
    path('complete_list', views.complete_list, name='complete_list'),  # 设计完善
    path('edit_complete', views.edit_complete, name='edit_complete'),  # 编辑设计完善
    path('upload_file', views.upload_file, name='upload_file'),  # 上传文件
    path('import_scenes', views.import_scenes, name='import_scenes'),  # 导入场景
    path('scenes_modeling', views.scenes_modeling, name='scenes_modeling'),  # 场景建模
    path('deliver_model_data', views.deliver_model_data, name='deliver_model_data'),  # 传递模型的场景信息
    path('model_list', views.model_list, name='model_list'),  # 模型列表

]


from flask import Blueprint, jsonify, request, request
from . import db
from .models import User, Product, Shop, Request
from sqlalchemy import join

main = Blueprint('main', __name__)

@main.route('/users', methods=['GET', 'POST'])
def users():
    if request.method == 'GET':
        user_list = User.query.all()
        users = []

        for user in user_list:
            users.append({'user_name' : user.user_name, 'rating' : user.address})

        return jsonify({'users' : users})
    else:
        user_data = request.get_json()

        new_user = User(user_name=user_data['user_name'], address=user_data['address'])

        db.session.add(new_user)
        db.session.commit()

        return 'Done', 201

@main.route('/shops', methods=['GET', 'POST'])
def shops():
    if request.method == 'GET':
        shop_list = Shop.query.all()
        shops = []

        for shop in shop_list:
            shops.append({'shop_name' : shop.shop_name, 'address' : shop.address, 'sells': shop.sells, 'id': shop.id})

        return jsonify({'shops' : shops})
    else:
        shop_data = request.get_json()

        new_shop = Shop(shop_name=shop_data['shop_name'], address=shop_data['address'], sells=shop_data['sells'])

        db.session.add(new_shop)
        db.session.commit()

        return 'Done', 201

@main.route('/requests', methods=['GET', 'POST'])
def requests():
    if request.method == 'GET':
        request_list = db.session.query(Shop, User, Product, Request).filter(Shop.id == Request.shop_id).filter(User.id == Request.user_id).filter(Product.id == Request.user_id)
        requests = []

        for re in request_list:
            requests.append({'shop_id': re.Shop.id, 'shop_name': re.Shop.shop_name, 'shop_address': re.Shop.address, 'user_id': re.User.id, 'user_name': re.User.user_name, 'user_address': re.User.address, 'product_id': re.Product.id, 'product_name': re.Product.product_name,'product_quantity': re.Product.quantity, 'request_quantity': re.Request.quantity})

        return jsonify({'requests' : requests})

    else:
        request_data = request.get_json()

        new_request = Request(shop_id=request_data['shop_id'], user_id=request_data['user_id'], product_id=request_data['product_id'], quantity=request_data['quantity'])

        db.session.add(new_request)
        db.session.commit()

        return 'Done', 201

@main.route('/products', methods=['GET', 'POST'])
def products():
    if request.method == 'GET':

        product_list =  db.session.query(Shop, Product).filter(Shop.id == Product.shop_id)
        products = []

        for product in product_list:
            products.append({'product_id': product.Product.id,  'product_name': product.Product.product_name, 'quantity': product.Product.quantity, 'shop_id': product.Shop.id, 'shop_name': product.Shop.shop_name, 'addres': product.Shop.address, 'sells': product.Shop.sells})

        return jsonify({'products' : products})
    else:
        product_data = request.get_json()

        new_product = Product(shop_id=product_data['shop_id'], product_name=product_data['product_name'], quantity=product_data['quantity'])

        db.session.add(new_product)
        db.session.commit()

        return 'Done', 201


from collections import Counter

class Solution:
    def findAnagrams(self, s: str, p: str) -> List[int]:

        L = 0
        result = []
        charMap = {}
        pCharMap = Counter(p)

        for R in range(len(s)):
            # Increment the charMap for each character in the window
            charMap[s[R]] = charMap.get(s[R], 0) + 1

            # Shrink the window until it matches the size of p. The code may seem weird but it makes sense
            # as we make L bigger ( by shifting the L pointer) we are subtracting byt a bigger L and therefore shrinks the window size
            while R - L + 1 > len(p):
                charMap[s[L]] -= 1
                if charMap[s[L]] == 0:
                    del charMap[s[L]]
                L += 1

            # Check if we have an anagram
            if pCharMap == charMap:
                result.append(L)

        return result


from django.urls import path

from deceases.views import symptoms_autocomplete, diagnostics, deceases_by_symptoms, symptom_tree, \
    doctor_create_update_decease, medical_records, decease_detail, decease_list, decease_autocomplete

urlpatterns = [
    path('medical-records/<int:patient_id>', medical_records, name='medical-records'),
    path('symptom_tree/', symptom_tree, name='symptom-tree'),
    path('diagnostics/', diagnostics, name='diagnostics'),
    path('detail/<int:pk>/', decease_detail, name='decease-detail'),
    path('list/', decease_list, name='decease-list'),
    path('autocomplete/', decease_autocomplete, name='decease-autocomplete'),
    path('symptoms-autocomplete/', symptoms_autocomplete, name='symptoms-autocomplete'),
    path('deceases-by-symptoms/', deceases_by_symptoms, name='deceases-by-symptoms'),
    path('doctor-create-update-patient-decease/', doctor_create_update_decease, name='doctor-create-patient-decease'),
    path('doctor-create-update-patient-decease/<int:pk>/', doctor_create_update_decease,
         name='doctor-update-patient-decease'),

]


import datetime

from bok_choy.web_app_test import WebAppTest

from analyticsclient.constants import demographic
from acceptance_tests import CoursePageTestsMixin
from acceptance_tests.pages import CourseEnrollmentActivityPage, CourseEnrollmentGeographyPage

_multiprocess_can_split_ = True

class CourseEnrollmentActivityTests(CoursePageTestsMixin, WebAppTest):
    def setUp(self):
        super(CourseEnrollmentActivityTests, self).setUp()
        self.page = CourseEnrollmentActivityPage(self.browser)
        self.course = self.api_client.courses(self.page.course_id)

    def get_enrollment_data(self):
        """
        Returns all historical enrollment data for enrollment count collection.
        """
        end_date = datetime.datetime.utcnow()
        end_date_string = end_date.strftime(self.api_client.DATE_FORMAT)
        return self.course.enrollment(start_date=None, end_date=end_date_string)

    def test_page(self):
        super(CourseEnrollmentActivityTests, self).test_page()
        self._test_enrollment_metrics_and_graph()
        self._test_enrollment_trend_table()

    def _get_data_update_message(self):
        current_enrollment = self.course.enrollment()[0]
        last_updated = datetime.datetime.strptime(current_enrollment['created'], self.api_datetime_format)
        return 'Enrollment activity data was last updated %(update_date)s at %(update_time)s UTC.' % \
               self.format_last_updated_date_and_time(last_updated)

    def _test_enrollment_metrics_and_graph(self):
        """ Verify the graph loads and that the metric tiles display the correct information. """

        enrollment_data = self.get_enrollment_data()
        current_enrollment = enrollment_data[-1]

        # Check values of summary boxes
        current_enrollment_count = current_enrollment['count']
        data_selector = 'data-stat-type=current_enrollment'
        self.assertSummaryPointValueEquals(data_selector, unicode(current_enrollment_count))
        self.assertSummaryTooltipEquals(data_selector, u'Students enrolled in the course.')

        # Check value of summary box for last week
        i = 7
        value = current_enrollment_count - enrollment_data[-(i + 1)]['count']
        data_selector = 'data-stat-type=enrollment_change_last_%s_days' % i
        self.assertSummaryPointValueEquals(data_selector, unicode(value))
        self.assertSummaryTooltipEquals(data_selector,
                                        u'The difference between the number of students enrolled at the end of the day yesterday and one week before.')

        # Verify *something* rendered where the graph should be. We cannot easily verify what rendered
        self.assertElementHasContent("[data-section=enrollment-basics] #enrollment-trend-view")

    def _test_enrollment_trend_table(self):
        """ Verify the information rendered in the table is correct. """

        enrollment_data = sorted(self.get_enrollment_data(), reverse=True, key=lambda item: item['date'])

        table_selector = 'div[data-role=enrollment-table] table'
        self.assertTableColumnHeadingsEqual(table_selector, ['Date', 'Total Enrollment'])

        rows = self.page.browser.find_elements_by_css_selector('%s tbody tr' % table_selector)
        self.assertGreater(len(rows), 0)

        for i, row in enumerate(rows):
            columns = row.find_elements_by_css_selector('td')
            enrollment = enrollment_data[i]
            expected_date = datetime.datetime.strptime(enrollment['date'], self.api_date_format).strftime(
                "%B %d, %Y").replace(' 0', ' ')
            expected = [expected_date, enrollment['count']]
            actual = [columns[0].text, int(columns[1].text)]
            self.assertListEqual(actual, expected)
            self.assertIn('text-right', columns[1].get_attribute('class'))

        # Verify CSV button has an href attribute
        selector = "a[data-role=enrollment-trend-csv]"
        self.assertValidHref(selector)

class CourseEnrollmentGeographyTests(CoursePageTestsMixin, WebAppTest):
    def setUp(self):
        super(CourseEnrollmentGeographyTests, self).setUp()
        self.page = CourseEnrollmentGeographyPage(self.browser)
        self.course = self.api_client.courses(self.page.course_id)
        self.enrollment_data = sorted(self.course.enrollment(demographic.LOCATION),
                                      key=lambda item: item['count'], reverse=True)

    def test_page(self):
        super(CourseEnrollmentGeographyTests, self).test_page()
        self._test_enrollment_country_map()
        self._test_enrollment_country_table()
        self._test_metrics()

    def _get_data_update_message(self):
        current_enrollment = self.course.enrollment(demographic.LOCATION)[0]
        last_updated = datetime.datetime.strptime(current_enrollment['created'], self.api_datetime_format)
        return 'Geographic student data was last updated %(update_date)s at %(update_time)s UTC.' % \
               self.format_last_updated_date_and_time(last_updated)

    def _test_enrollment_country_map(self):
        """ Verify the geolocation map is loaded. """

        map_selector = "div[data-view=world-map]"

        # Ensure the map is loaded via AJAX
        self.fulfill_loading_promise(map_selector)

        # make sure the map section is present
        element = self.page.q(css=map_selector)
        self.assertTrue(element.present)

        # make sure that the map is present
        element = self.page.q(css=map_selector + " svg[class=datamap]")
        self.assertTrue(element.present)

        # make sure the legend is present
        element = self.page.q(css=map_selector + " svg[class=datamaps-legend]")
        self.assertTrue(element.present)

    def _test_enrollment_country_table(self):
        """ Verify the geolocation enrollment table is loaded. """

        table_section_selector = "div[data-role=enrollment-location-table]"

        # Ensure the table is loaded via AJAX
        self.fulfill_loading_promise(table_section_selector)

        # make sure the map section is present
        element = self.page.q(css=table_section_selector)
        self.assertTrue(element.present)

        # make sure the table is present
        table_selector = table_section_selector + " table"
        element = self.page.q(css=table_selector)
        self.assertTrue(element.present)

        # check the headings
        self.assertTableColumnHeadingsEqual(table_selector, ['Country', 'Percent', 'Total Enrollment'])

        # Verify CSV button has an href attribute
        selector = "a[data-role=enrollment-location-csv]"
        self.assertValidHref(selector)

        # Check the results of the table
        rows = self.page.browser.find_elements_by_css_selector('%s tbody tr' % table_selector)
        self.assertGreater(len(rows), 0)

        sum_count = float(sum([datum['count'] for datum in self.enrollment_data]))

        for i, row in enumerate(rows):
            columns = row.find_elements_by_css_selector('td')
            enrollment = self.enrollment_data[i]
            expected_percent = enrollment['count'] / sum_count * 100
            expected_percent_display = '{:.1f}%'.format(expected_percent) if expected_percent >= 1.0 else '< 1%'
            expected = [enrollment['country']['name'], expected_percent_display, enrollment['count']]
            actual = [columns[0].text, columns[1].text, int(columns[2].text)]
            self.assertListEqual(actual, expected)
            self.assertIn('text-right', columns[1].get_attribute('class'))

    def _test_metrics(self):
        """ Verify the metrics tiles display the correct information. """

        enrollment_data = [datum for datum in self.enrollment_data if datum['country']['name'] != 'UNKNOWN']
        self.assertSummaryPointValueEquals('data-stat-type=num-countries', unicode(len(enrollment_data)))

        for i in range(0, 3):
            country = enrollment_data[i]['country']['name']
            self.assertSummaryPointValueEquals('data-stat-type=top-country-{0}'.format((i + 1)), country)


from __future__ import print_function
from pyspark.sql import functions as F

def cal_psi(df1, df2, col_name, n_bin = 20, bin_type ="quantiles", verbose =False):
    """

    :param df1: dataframe
    :param df2: dataframe
    :param col_name: column to calculate psi
    :param n_bin: number bin
    :param bin_type: is "quantiles" or "bins"
    :param verbose: print explaination
    :return: psi
    """
    if bin_type not in ("quantiles", "bins"):
        print("need bin_type = quantiles or bins")
        raise ValueError
        return
    df_union = df1.select(col_name)\
        .unionByName(df2.select(col_name))
    if bin_type == "quantiles":
        delta = 1.0/float(n_bin)
        prob_list = [0.0 + delta * i for i in range(n_bin)]
        point_list = df_union.stat.approxQuantile(col_name, prob_list, 0.001)
    if bin_type == "bins":
        min_max = df_union\
            .select(F.min(col_name).alias("min"), F.max(col_name).alias("max"))\
            .collect()[0]
        min_value = min_max["min"]
        max_value = min_max["max"]
        delta = (max_value - min_value)/ float(n_bin)
        point_list = [min_value + delta * i for i in range(n_bin)]
    bin_expr = ""
    for i in range(len(point_list)):
        if i ==0:
            # first item only compare <
            col_value = "when {0} < {1} then {2} \n "\
                .format(col_name, point_list[i + 1], i + 1)
        elif (i == len(point_list) -1):
            # final item: only compare >=
            col_value = "when {0} > {1} then {2} \n "\
                        .format(col_name, point_list[i], i + 1)
        else:
            col_value = "when {0} between {1} and {2} then {3} \n "\
                .format(col_name, point_list[i], point_list[i + 1], i + 1)
        bin_expr += col_value
    # Null value bin =0, NAN value bin = max_bin.
    bin_expr = "case \n " + bin_expr + " else 0 \n end"
    if verbose:
        print("bin_expr",bin_expr)
        print("point_list",point_list)
    df1_count = df1.count()
    df1_agg = df1.withColumn("bin", F.expr(bin_expr))\
        .groupBy("bin")\
        .agg(
            (F.count("*")/df1_count).alias("count_per1"))
    df2_count = df2.count()
    df2_agg = df2.withColumn("bin", F.expr(bin_expr))\
        .groupBy("bin")\
        .agg(
            (F.count("*")/df2_count).alias("count_per2"))
    df_psi = df1_agg\
        .join(df2_agg, "bin", "outer")\
        .withColumn("psi", (F.expr("count_per1 - count_per2") *F.log(F.expr("count_per1/count_per2"))))
    psi = df_psi.select(F.sum("psi").alias("psi")).collect()[0]["psi"]
    return psi


"""
Module initialize the unfolding.

@title: initialisation
@author: Valentin Louf <valentin.louf@bom.gov.au>
@institutions: Monash University and the Australian Bureau of Meteorology
@date: 16/03/2021

.. autosummary::
    :toctree: generated/

    find_last_good_vel
    flipud
    first_pass
    initialize_unfolding
"""
# Other Libraries
import numpy as np

from numba import jit, jit_module
from numba import uint32, int64, float64

# Custom
from .continuity import take_decision, unfold, is_good_velocity

def find_last_good_vel(j, n, azipos, vflag, nfilter):
    """
    Looking for the last good (i.e. processed) velocity in a slice.

    Parameters:
    ===========
    j: int
        Position in dimension 1
    n: int
        Position in dimension 2
    azipos: array
        Array of available position in dimension 1
    vflag: ndarray <azimuth, range>
        Flag array (-3: missing, 0: unprocessed, 1: processed, 2: processed and unfolded)
    nfilter: int
        Size of filter.

    Returns:
    ========
        idx_ref: int
            Last valid value in slice.
    """
    i = 0
    while i < nfilter:
        i += 1
        idx_ref = j - i
        idx = azipos[idx_ref]
        vflag_ref = vflag[idx, n]
        if vflag_ref > 0:
            return idx_ref
    return -999

def flipud(arr):
    """
    Numpy's flipud function is not supported by numba for some reasons...
    So here it is.
    """
    return arr[::-1, :]  # Soooo complex!

def first_pass(azi_start_pos, velocity, final_vel, vflag, vnyquist, delta_vmax, nfilter=3):
    """
    First pass: continuity check along the azimuth, starting at azi_start_pos.

    Parameters:
    ===========
    azi_start_pos: int
        Starting position in array alongside dimension 0.
    velocity: ndarray <azimuth, range>
        Velocity to dealias
    final_vel: ndarray <azimuth, range>
        Result array
    vflag: ndarray <azimuth, range>
        Flag array (-3: missing, 0: unprocessed, 1: processed, 2: processed and unfolded)
    vnyquist: float
        Nyquist velocity.
    vshift: float
        Shift expected when folding, i.e. 2*vnyquist (in general)
    delta_vmax: float
        Maximum velocity difference tolerated between two contiguous gates.
    nfilter: int
        Size of filtering window

    Returns:
    ========
    final_vel: array <azimuth, range>
        Unfolded velocities.
    flag_vel: array <azimuth, range>
        Flag array for velocity processing (0: unprocessed, 1:processed, 2:unfolded, -3: missing)
    """
    nazi, ngate = velocity.shape
    azipos = np.zeros((2 * nazi), dtype=uint32)
    azipos[:nazi] = np.arange(nazi)
    azipos[nazi:] = np.arange(nazi)

    for mypass in range(2):
        if mypass == 1:
            velocity = flipud(velocity)
            final_vel = flipud(final_vel)
            vflag = flipud(vflag)
            azi_start_pos = nazi - azi_start_pos - 1

        for j in range(azi_start_pos, azi_start_pos + nazi // 2):
            for n in range(ngate):
                # Build slice for comparison
                j_idx = np.arange(j + 1, j + nfilter + 1)
                j_idx[j_idx >= nazi] -= nazi

                idx_selected = vflag[j_idx, n]
                if np.all((idx_selected == -3)):
                    # All values missing in slice.
                    continue
                # Selection of velocity to dealias (j_idx is an array)
                v_selected = velocity[j_idx, n]

                # Searching reference
                idx_ref = j
                # Looking for last processed value for reference, within a nfilter distance.
                if vflag[azipos[j], n] <= 0:
                    idx_ref = find_last_good_vel(j, n, azipos, vflag, nfilter)
                    if idx_ref == -999:
                        continue

                # Reference velocity
                vref = final_vel[azipos[int(idx_ref)], n]

                # Dealiasing slice
                for k in range(len(v_selected)):
                    if idx_selected[k] == -3:
                        continue

                    vk = v_selected[k]
                    dv1 = np.abs(vk - vref)
                    if dv1 < delta_vmax:
                        # No need to dealias
                        final_vel[j_idx[k], n] = vk
                        vflag[j_idx[k], n] = 1
                    else:
                        vk_unfld = unfold(vref, vk, vnyquist)
                        dvk = np.abs(vk_unfld - vref)

                        if dvk < delta_vmax:
                            final_vel[j_idx[k], n] = vk_unfld
                            vflag[j_idx[k], n] = 2

    velocity = flipud(velocity)
    final_vel = flipud(final_vel)
    vflag = flipud(vflag)

    return final_vel, vflag

def initialize_unfolding(azi_start_pos, azi_end_pos, vel, flag_vel, vnyq=13.3):
    """
    Initialize the unfolding procedure and unfold the reference radials..

    Parameters:
    ===========

    Returns:
    ========
    final_vel: array <azimuth, range>
        Unfolded velocities.
    flag_vel: array <azimuth, range>
        Flag array for velocity processing (0: unprocessed, 1:processed, 2:unfolded, -3: missing)
    """
    # Initialize stuff.
    maxazi = vel.shape[0]
    final_vel = np.zeros(vel.shape, dtype=float64)

    iter_radials_init = np.array([azi_start_pos - 1, azi_start_pos, azi_start_pos + 1])
    iter_radials_last = np.array([azi_end_pos - 1, azi_end_pos, azi_end_pos + 1])

    for iter_radials in [iter_radials_init, iter_radials_last]:
        iter_radials[iter_radials >= maxazi] -= maxazi
        iter_radials[iter_radials < 0] += maxazi

        # Magic happens.
        is_bad = 0
        for pos_good in iter_radials:
            myvel = vel[pos_good, :]
            if np.sum(np.abs(myvel) >= 0.6 * vnyq) > 3:
                is_bad += 1
                continue

            for ngate in range(3, len(myvel) - 3):
                velref0 = np.nanmedian(myvel[ngate - 3 : ngate])
                velref1 = np.nanmedian(myvel[ngate + 1 : ngate + 4])
                decision = take_decision(velref0, velref1, vnyq, alpha=0.4)
                if decision != 1:
                    continue

                if np.isnan(velref0):
                    velref = velref1
                elif np.isnan(velref1):
                    velref = velref0
                else:
                    velref = (velref0 + velref1) / 2

                decision = take_decision(velref, myvel[ngate], vnyq, alpha=0.4)
                if decision == 0:
                    continue
                elif decision == 1:
                    final_vel[pos_good, ngate] = myvel[ngate]
                    flag_vel[pos_good, ngate] = 1
                elif decision == 2:
                    vtrue = unfold(myvel[ngate - 1], myvel[ngate], vnyq)
                    if is_good_velocity(myvel[ngate - 1], vtrue, vnyq, alpha=0.4):
                        final_vel[pos_good, ngate] = vtrue
                        flag_vel[pos_good, ngate] = 2

        if is_bad > len(iter_radials) - 1:
            for pos_good in iter_radials:
                myvel = vel[pos_good, :]

                for ngate in range(3, len(myvel) - 3):
                    if np.abs(myvel[ngate]) >= 0.5 * vnyq:
                        continue

                    velref0 = np.nanmedian(myvel[ngate - 3 : ngate])
                    velref1 = np.nanmedian(myvel[ngate + 1 : ngate + 4])
                    decision = take_decision(velref0, velref1, vnyq, alpha=0.4)
                    if decision != 1:
                        continue

                    if np.isnan(velref0):
                        velref = velref1
                    elif np.isnan(velref1):
                        velref = velref0
                    else:
                        velref = (velref0 + velref1) / 2

                    decision = take_decision(velref, myvel[ngate], vnyq, alpha=0.4)
                    if decision == 0:
                        continue
                    elif decision == 1:
                        final_vel[pos_good, ngate] = myvel[ngate]
                        flag_vel[pos_good, ngate] = 1
                    elif decision == 2:
                        vtrue = unfold(myvel[ngate - 1], myvel[ngate], vnyq)
                        if is_good_velocity(myvel[ngate - 1], vtrue, vnyq, alpha=0.4):
                            final_vel[pos_good, ngate] = vtrue
                            flag_vel[pos_good, ngate] = 2

    return final_vel, flag_vel

jit_module(nopython=True, error_model="numpy")