import torch

import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from .embedding import Optcodes
from hydra.utils import instantiate

from core.utils.ray_utils import kp_to_valid_rays
from core.utils.skeleton_utils import (
    Skeleton,
    SMPLSkeleton,
    get_skel_profile_from_rest_pose,
    calculate_kinematic,
)

from einops import rearrange
from core.networks.camcal import CamCal
from typing import Optional, Callable, Mapping, Any, List, Union
from omegaconf import DictConfig

def batchify_fn(
    fn: Callable,
    fn_inputs: Mapping[str, Any],
    N_total: int,
    render_normal: bool = False,
    chunk: int = 4096
):
    """ Break evaluation into batches to avoid OOM
    """
    all_ret = {}

    for i in range(0, N_total, chunk):
        batch_inputs = {k: fn_inputs[k][i:i+chunk] if torch.is_tensor(fn_inputs[k]) else fn_inputs[k]
                        for k in fn_inputs}
        ret = fn(batch_inputs, render_normal=render_normal)
        if ret is None:
            continue

        for k in ret:
            if k not in all_ret:
                all_ret[k] = []
            all_ret[k].append(ret[k])

    all_ret = {k : torch.cat(all_ret[k], 0) for k in all_ret}
    return all_ret

def merge_encodings(
    encoded: Mapping[str, torch.Tensor],
    encoded_is: Mapping[str, torch.Tensor],
    sorted_idxs: torch.Tensor,
    N_rays: int,
    N_total_samples: int,
    inplace: bool = True
):
    """
    merge coarse and fine encodings.
    Parameters
    ----------
    encoded: dictionary of coarse encodings
    encoded_is: dictionary of fine encodings
    sorted_idxs: define how the [encoded, encoded_is] are sorted
    """
    if encoded_is is None:
        return encoded
    if encoded is None:
        return encoded_is
    gather_idxs = torch.arange(N_rays * (N_total_samples)).view(N_rays, -1)
    gather_idxs = torch.gather(gather_idxs, 1, sorted_idxs)

    merged = encoded if inplace else {}

    for k in encoded.keys():
        #if not k.startswith(('pts', 'blend')) and k not in ['g', 'graph_feat', 'bone_logit']:
        if k in ['valid_idxs', 'rw'] or encoded[k] is None or isinstance(encoded[k], dict):
            continue
        if k in encoded and k not in encoded_is:
            merged[k] = encoded[k]
            continue
        merged[k] = merge_samples(encoded[k], encoded_is[k], gather_idxs, N_total_samples)

    # need special treatment here to preserve the computation graph.
    # (otherwise we can just re-encode everything again, but that takes extra computes)
    if 'pts' in encoded and encoded['pts'] is not None:
        if not inplace:
            merged['pts'] = encoded['pts']

        merged['pts_is'] = encoded_is['pts']
        merged['gather_idxs'] = gather_idxs

        merged['pts_sorted'] = merge_samples(encoded['pts'], encoded_is['pts'],
                                                gather_idxs, N_total_samples)

    return merged

def merge_samples(
    x: torch.Tensor,
    x_is: torch.Tensor,
    gather_idxs: torch.Tensor,
    N_total_samples: int
):
    """
    merge coarse and fine samples.
    Parameteters
    ------------
    x: coarse samples of shape (N_rays, N_coarse, -1)
    x_is: importance samples of shape (N_rays, N_fine, -1)
    gather_idx: define how the [x, x_is] are sorted
    """
    if x is None or x.shape[-1] == 0:
        return None
    N_rays = x.shape[0]
    if x.shape[0] != x_is.shape[0]:
        return torch.cat([x, x_is], dim=0)
    x_is = torch.cat([x, x_is], dim=1)
    sh = x_is.shape
    feat_size = np.prod(sh[2:])
    if x_is.shape[1] != N_total_samples:
        # extra signal: may not have the same shape as the standard ray samples
        return x_is
    x_is = x_is.view(-1, feat_size)[gather_idxs, :]
    x_is = x_is.view(N_rays, N_total_samples, *sh[2:])

    return x_is

class ANeRF(nn.Module):

    def __init__(
        self,
        D: int,
        W: int,
        view_W: int,
        pts_embedder: DictConfig,
        pts_posi_enc: DictConfig,
        view_embedder: DictConfig,
        view_posi_enc: DictConfig,
        raycaster: DictConfig,
        skel_type: Skeleton,
        rest_pose: np.ndarray,
        skips: List[int] = [4],
        pred_sdf: bool = False,
        use_framecodes: bool = False,
        framecode_ch: int = 16,
        n_framecodes: int = 0,
        density_noise_std: float = 1.0,
        bkgd_net: Optional[DictConfig] = None,
        cam_cal: Optional[DictConfig] = None,
        pose_opt: Optional[DictConfig] = None,
        **kwargs,
    ):
        '''
        Parameters
        ----------
        D: int, depth of the MLP
        W: int, width of the MLP
        view_W: int, width of the view MLP
        pts_embedder: embedder module config, to encode 3d points w.r.t. body keypoints and poses
        pts_posi_enc: positional encoding config for the pts embedding
        view_embedder: embedder config, to encode view vectors w.r.t. body info
        view_posi_enc: positional encoding config for the view embedding
        raycast: ray casting module config
        skel_type: skeleton, define the details of the skeleton
        skips: list, layers to do skip connection
        pred_sdf: Bool, pred_sdf=True then predicts SDF value
        use_framecodes: Bool, to use framecodes for each frame
        framecode_ch: int, size of the framecode
        n_framecodes: int, number of framecodes
        density_noise_std: float, noise to apply on density during training time
        '''
        super(ANeRF, self).__init__()

        self.D = D
        self.W = W
        self.view_W = view_W
        self.skips = skips
        self.pred_sdf = pred_sdf
        self.use_framecodes = use_framecodes
        self.framecode_ch = framecode_ch
        self.n_framecodes = n_framecodes
        self.skel_type = skel_type
        self.density_noise_std = density_noise_std

        # initialize skeleton settings
        self.init_skeleton(skel_type, rest_pose)
        # instantiate embedder and network
        self.init_embedder(
            pts_embedder=pts_embedder,
            pts_posi_enc=pts_posi_enc,
            view_embedder=view_embedder,
            view_posi_enc=view_posi_enc,
            **kwargs
        )
        self.init_raycast(raycaster, **kwargs)
        self.init_density_net()
        self.init_radiance_net()
        self.init_bkgd_net(bkgd_net)

        self.cam_cal = None
        if cam_cal is not None:
            self.cam_cal = CamCal(**cam_cal)

        self.pose_opt = None
        if pose_opt is not None:
            self.pose_opt = instantiate(
                pose_opt,
                rest_pose=self.rest_pose,
                skel_type=self.skel_type,
                **kwargs
            )

    @property
    def dnet_input(self):
        return self.input_ch

    @property
    def vnet_input(self):
        if self.use_framecodes:
            return self.input_ch_view + self.framecode_ch + self.W
        return self.input_ch_view + self.W

    def init_skeleton(self, skel_type: Skeleton, rest_pose: np.ndarray):
        self.register_buffer('rest_pose', torch.tensor(rest_pose))
        self.skel_profile = get_skel_profile_from_rest_pose(rest_pose, skel_type)

    def init_embedder(
        self,
        pts_embedder: DictConfig,
        pts_posi_enc: DictConfig,
        view_embedder: DictConfig,
        view_posi_enc: DictConfig,
        **kwargs,
    ):
        N_joints = len(self.skel_type.joint_names)
        # initialize points transformation
        self.pts_embedder = instantiate(pts_embedder, N_joints=N_joints, N_dims=3)
        pts_dims, pose_dims = self.pts_embedder.dims

        # initialize positional encoding for points
        self.pts_dims = pts_dims
        self.pose_dims = pose_dims
        self.pts_posi_enc = instantiate(pts_posi_enc, input_dims=pts_dims)

        # initialize view transformation
        self.view_embedder = instantiate(view_embedder, N_joints=N_joints, N_dims=3)

        # initialize positional encoding for views (rays)
        self.view_dims = view_dims= self.view_embedder.dims
        self.view_posi_enc = instantiate(view_posi_enc, input_dims=view_dims, dist_inputs=True)

        self.input_ch = self.pts_posi_enc.dims + pose_dims
        self.input_ch_view = self.view_posi_enc.dims

    def init_raycast(
        self,
        raycast: DictConfig,
        **kwargs,
    ):
        self.raycast = instantiate(raycast, **kwargs, _recursive_=False)

    def init_density_net(self):

        W, D = self.W, self.D

        layers = [nn.Linear(self.dnet_input, W)]

        for i in range(D-1):
            if i not in self.skips:
                layers += [nn.Linear(W, W)]
            else:
                layers += [nn.Linear(W + self.dnet_input, W)]

        self.pts_linears = nn.ModuleList(layers)
        self.sigma_linear = nn.Linear(W, 1)

        if self.pred_sdf:
            self.register_parameter('B',
                nn.Parameter(torch.tensor(0.1), requires_grad=True)
            )

    def init_radiance_net(self):

        W, view_W = self.W, self.view_W

        # Note: legacy code, don't really need nn.ModuleList
        self.views_linears = nn.ModuleList([nn.Linear(self.vnet_input, view_W)])
        self.feature_linear = nn.Linear(W, view_W * 2)
        self.rgb_linear = nn.Linear(view_W, 3)

        if self.use_framecodes:
            self.framecodes = Optcodes(self.n_framecodes, self.framecode_ch)

    def init_bkgd_net(self, bkgd_net: DictConfig):
        if bkgd_net is None:
            self.bkgd_net = None
            return
        self.bkgd_net = instantiate(
            bkgd_net,
            n_framecodes=self.n_framecodes,
            framecode_ch=self.framecode_ch
        )

    def forward(self, *args, forward_type: str = 'rays', **kwargs):
        """ Forward function.

        Args:
            forward_type (str): ['rays', 'render', 'geometry']

        Returns:
            Dictionary of output tensors.
        """
        if forward_type == 'rays':
            return self.forward_rays(*args, **kwargs)
        elif forward_type == 'render':
            return self.forward_render(*args, **kwargs)
        elif forward_type =='geometry':
            return self.forward_geometry(*args, **kwargs)
        else:
            raise NotImplementedError(f'Unknown forward type {forward_type}')

    def forward_rays(self, batch: Mapping[str, Any], pose_opt: bool = False, **kwargs):

        if 'rays_o' not in batch and 'rays_d' not in batch:
            raise NotImplementedError(
                'Rays are not provided as input. '
                'For rendering image with automatically detected rays, set forward_type="render"'
            )

        # Step 1. cast ray
        bgs = batch.get('bgs', None) # background color
        sample_info = self.raycast(batch)
        pts, z_vals = sample_info['pts'], sample_info['z_vals']

        # Step 2. model evaluation
        # TODO: do we need a get_nerf_inputs function?
        network_inputs = self.get_network_inputs(batch, pts, z_vals)
        if pose_opt and self.pose_opt is not None:
            # TODO: is this a good way?
            network_inputs = self.pose_opt(
                network_inputs=network_inputs,
                kp3d=network_inputs['kp3d'],
                bones=network_inputs['bones'],
                kp_idxs=network_inputs['real_kp_idx'],
                N_unique=network_inputs['N_unique'],
            )

        raw, encoded = self.evaluate_pts(network_inputs, coarse=True)

        if raw is None:
            return self._empty_outputs(batch)

        bgnet_ret = None
        if self.training and self.bkgd_net is not None:
            bgnet_ret = self.bkgd_net(batch)
            bgs = bgnet_ret['bgs']

        # Step 3. coarse rendering
        rendered = self.raw2outputs(
            raw,
            z_vals,
            batch['rays_d'],
            bgs=bgs,
            encoded=encoded,
            batch=batch,
        )

        if self.raycast.N_importance == 0:
            return {'rendered': rendered, 'encoded': encoded}

        # Step 4. importance sampling (if importance sampling enabled)
        rendered_coarse = rendered
        encoded_coarse = encoded

        weights = rendered_coarse['weights']
        is_sample_info = self.raycast(
            batch,
            pts=pts,
            z_vals=z_vals,
            weights=weights,
            importance=True,
        )
        pts_is = is_sample_info['pts'] # only importance samples
        z_vals = is_sample_info['z_vals'] # include both coarse and importance samples

        # Step 5. model evaluation (if importance sampling enabled)
        network_inputs = self.get_network_inputs(batch, pts_is, is_sample_info['z_vals_is'])
        raw_is, encoded_is = self.evaluate_pts(network_inputs, encoded_coarse=encoded_coarse)

        # Step 6. merge coarse and importance prediction for rendering
        sorted_idxs = is_sample_info['sorted_idxs']
        N_rays = len(batch['rays_o'])
        N_total_samples = pts.shape[1] + pts_is.shape[1]
        encoded_is = merge_encodings(encoded_coarse, encoded_is, sorted_idxs, N_rays, N_total_samples)

        # Step 7. fine rendering
        if raw_is is not None:
            raw = merge_encodings({'raw': raw}, {'raw': raw_is}, sorted_idxs, N_rays, N_total_samples)['raw']
            rendered = self.raw2outputs(
                raw,
                z_vals,
                batch['rays_d'],
                bgs=bgs,
                encoded=encoded_is,
                batch=batch,
            )
        else:
            rendered = rendered_coarse

        if bgnet_ret is not None:
            rendered.update(**bgnet_ret)

        return self.collect_outputs(rendered, rendered_coarse, encoded_is, encoded_coarse)

    def forward_render(
        self,
        render_data: Mapping[str, Any],
        render_factor: float = 0,
        raychunk: int = 1024*10,
        render_normal: bool = False,
        **kwargs
    ):
        """ Render images with automatically detected rays from camera parameters.

        Args:
            render_data (Mapping[str, Any]): dictionary of rendering data, including skeleton structure and camera info
            render_factor (float): downsample factor for rendering
            raychunk (int): chunk size for ray casting
            render_normal (bool): whether to render surface normal

        Returns:
            Dictionary of rendered outputs.
        """

        assert 'bones' in render_data, 'needs know the pose/bones (bones) parameter during rendering'
        if 'kp3d' not in render_data:
            kp3d, skts = calculate_kinematic(
                self.rest_pose,
                render_data['bones'],
                render_data.get('root_locs', None),
            )
            render_data['kp3d'] = kp3d
            render_data['skts'] = skts

        H, W, focals = render_data['hwf']
        kp3d = render_data['kp3d']
        cam_poses = render_data['c2ws'] # camera-to-world
        centers = render_data.get('center', None)

        if render_factor != 0:
            # change camera setting
            H, W = H // render_factor, W // render_factor
            focals = focals / render_factor
            if centers is not None:
                centers = centers / render_factor
            bgs = render_data['bgs']
            N, _H, _W, C = bgs.shape

            bgs = rearrange(bgs, 'n h w c -> n c h w')
            bgs = F.interpolate(bgs, scale_factor=1 / render_factor, mode='bilinear', align_corners=False)
            bgs = rearrange(bgs, 'n c h w -> n h w c')
            render_data.update(bgs=bgs)

        if len(cam_poses) != len(kp3d):
            assert len(kp3d) == 1 or len(cam_poses) == 1, \
                   f'Number of body poses should either match or can be broadcasted to number of camera poses. ' \
                   f'Got {len(kp3d)} and {len(cam_poses)}'

        rgb_imgs = []
        disp_imgs = []
        alpha_imgs = []

        for i in range(len(cam_poses)):

            # Step 1. find valid rays for this camera
            center = centers[i:i+1] if centers is not None else None
            rays, valid_idxs, cyls, _ = kp_to_valid_rays(
                cam_poses[i:i+1],
                H[i:i+1].cpu().numpy(),
                W[i:i+1].cpu().numpy(),
                focals[i:i+1].cpu().numpy(),
                kps=kp3d[i:i+1],
                centers=center,
            )
            rays_o, rays_d = rays[0]
            valid_idxs = valid_idxs[0]

            # initialize images
            bg = render_data['bgs'][i].cpu()

            if valid_idxs is None or (len(valid_idxs) == 0):
                rgb_imgs.append(bg.clone())
                disp_imgs.append(torch.zeros_like(bg[..., 0]).cpu())
                continue

            # flatten to (H * W, 3)
            rgb_img = bg.clone().flatten(end_dim=1)
            disp_img = torch.zeros_like(bg[..., 0]).flatten(end_dim=1)

            # Step 2. turn them into the format that forward_rays takes
            render_data.update(cyls=cyls)
            batch = self.to_ray_inputs(rays_o, rays_d, render_data, valid_idxs, index=i)

            # Step 3. forward
            if not self.pred_sdf:
                with torch.no_grad():
                    preds = batchify_fn(self.forward_rays, batch, N_total=len(rays_o), chunk=raychunk)
            else:
                preds = batchify_fn(self.forward_rays, batch, N_total=len(rays_o),
                                    chunk=raychunk, render_normal=render_normal)
                preds = {k: v.detach() for k, v in preds.items()}

            # put the rendered values into images
            pred_rgb = preds['rgb_map'].detach()
            rgb_img[valid_idxs] = pred_rgb.cpu()
            rgb_img = rgb_img.reshape(H[i], W[i], 3)

            #pred_disp = preds['disp_map'].detach()
            # TODO: hack, fix this
            pred_disp = preds['acc_map']
            disp_img[valid_idxs] = pred_disp.cpu()
            disp_img = disp_img.reshape(H[i], W[i], 1)

            rgb_imgs.append(rgb_img)
            disp_imgs.append(disp_img)

        rgb_imgs = torch.stack(rgb_imgs, dim=0)
        disp_imgs = torch.stack(disp_imgs, dim=0)

        return {
            'rgb_imgs': rgb_imgs,
            'disp_imgs': disp_imgs
        }

    def get_network_inputs(
        self,
        batch: Mapping[str, Any],
        pts: torch.Tensor,
        z_vals: torch.Tensor,
        keys_from_batch: List[str] = [
            'kp3d', 'skts', 'bones', 'cam_idxs',
            'rays_o', 'rays_d', 'temp_kp', 'temp_bone',
            'temp_skt', 'real_kp_idx', 'real_cam_idx'],
        **kwargs
    ):
        """ Collect network inputs from the batch.

        Args:
            batch (Mapping[str, Any]): batch data
            pts (torch.Tensor): 3D points
            z_vals (torch.Tensor): z values of the points
            keys_from_batch (List[str]): specifying which keys to collect from the batch

        Returns:
            Dictionary containing network inputs.
        """
        ret = {k: batch[k] for k in keys_from_batch if k in batch}
        ret['pts'] = pts
        ret['N_unique'] = batch.get('N_unique', 1)

        if self.cam_cal is not None:
            ret = self.cam_cal(ret, z_vals)

        return ret

    def to_ray_inputs(
        self,
        rays_o: torch.Tensor,
        rays_d: torch.Tensor,
        render_data: Mapping[str, Any],
        valid_idxs: torch.Tensor,
        index: int,
        keys_for_batch: List[str] = ['kp3d', 'skts', 'bones', 'cyls', 'cam_idxs', 'bgs'],
        pixel_data: List[str] = ['bgs'],
    ):
        """ Turning ray data into the format that forward_rays takes.

        Args:
            rays_o (torch.Tensor): ray origins
            rays_d (torch.Tensor): ray directions
            render_data (Mapping[str, Any]): dictionary of rendering data
            valid_idxs (torch.Tensor): valid ray indices
            index (int): index of the current camera
            keys_for_batch (List[str]): specifying which keys to collect from the batch
            pixel_data (List[str]): list of keys that have data in (H, W, C) format

        Returns:
            Dictionary containing inputs for the forward_ray function
        """

        device = rays_o.device
        ray_inputs = {
            'rays_o': rays_o,
            'rays_d': rays_d,
            'N_unique': 1, # always one pose at a time
        }

        N_rays = len(rays_o)
        # TODO: hacky to define it here..

        for k in keys_for_batch:
            v = render_data[k]

            if len(v) > 1:
                v = v[index:index+1]

            if k not in pixel_data:
                sh = v.shape[1:]
                v = v.expand(N_rays, *sh)
            else:
                v = rearrange(v, 'b h w c -> (b h w) c')[valid_idxs.to(v.device)].to(device)
            ray_inputs[k] = v
        return ray_inputs

    def evaluate_pts(self, inputs: Mapping[str, Any], geometry_only: bool = False, geom_extra_rets: List[str] = [], **kwargs):
        """ Evaluate the NeRF function at the given inputs/points.
        """

        if self.pred_sdf:
            inputs['pts'].requires_grad = True

        # Step 1. encode all pts feature
        density_inputs, encoded_pts = self.encode_pts(inputs)
        if density_inputs is None:
            # terminate because not valid points
            return None, None

        # Step 2. density prediction
        sigma, density_feature = self.inference_sigma(density_inputs)

        if self.pred_sdf:
            surface_gradient = torch.autograd.grad(
                outputs=sigma,
                inputs=inputs['pts'],
                grad_outputs=torch.ones_like(sigma),
                create_graph=self.training,
            )[0]
            encoded_pts.update(surface_gradient=surface_gradient)

        density = self.to_density(sigma)

        if geometry_only:
            geom_ret = {'density': density, 'sigma': sigma}
            if 'valid_idxs' in encoded_pts:
                geom_ret['valid_idxs'] = encoded_pts['valid_idxs']
            if 'surface_gradient' in encoded_pts:
                geom_ret['surface_gradient'] = encoded_pts['surface_gradient']
            for k in geom_extra_rets:
                if k in encoded_pts:
                    geom_ret[k] = encoded_pts[k]
                if k == 'density_feature':
                    geom_ret[k] = density_feature
            return None, geom_ret

        # Step 3. encode all ray feature
        view_inputs, encoded_views = self.encode_views(
            inputs,
            refs=encoded_pts['pts_t'],
            encoded_pts=encoded_pts
        )

        # Step 4. rgb prediction
        rgb = self.inference_rgb(view_inputs, density_feature)

        # Step 5: create final outputs
        shape = inputs['pts'].shape[:2] # (N_rays, N_samples)
        outputs = self.fill_prediction(torch.cat([rgb, density], dim=-1), shape, encoded_pts)

        return outputs, self.collect_encoded(encoded_pts, encoded_views)

    def encode_pts(self, inputs: Mapping[str, Any]):
        """ Encode input points into encoding needed for the NeRF model.

        Args:
            inputs (Mapping[str, Any]): dictionary of inputs

        Returns:
            density_inputs (torch.Tensor): density network inputs
            encoded (Mapping[str, torch.Tensor]): detail encoding of the points
        """

        N_rays, N_samples = inputs['pts'].shape[:2]

        encoded = self.pts_embedder(**inputs)
        v_pe = self.pts_posi_enc(encoded['v'])[0]
        r = encoded['r'].reshape(N_rays, N_samples, self.pose_dims)

        # apply positional encoding (PE)
        # pe_fn returns a tuple: (encoded outputs, cutoff weights)
        density_inputs = torch.cat([v_pe, r], dim=-1).flatten(end_dim=1)

        return density_inputs, encoded

    def encode_views(self, inputs: Mapping[str, Any], refs: torch.Tensor, encoded_pts: Mapping[str, torch.Tensor]):
        """ Encode input views into encoding needed for the NeRF model.

        Args:
            inputs (Mapping[str, Any]): dictionary of inputs
            refs (torch.Tensor): reference tensor for expanding rays
            encoded_pts (Mapping[str, Any]): point encoding that could be useful for encoding view

        Returns:
            view_inputs (torch.Tensor): view network inputs
            encoded (Mapping[str, Any]): detail encoding of the views
        """

        N_rays, N_samples = refs.shape[:2]
        encoded = self.view_embedder(**inputs, refs=refs)
        d = encoded['d'].reshape(N_rays, N_samples, self.view_dims)

        # apply positional encoding (PE)
        d_pe = self.view_posi_enc(d, dists=encoded_pts.get('v', None))[0]

        view_inputs = d_pe
        if self.use_framecodes:
            N_rays, N_samples = refs.shape[:2]
            # expand from (N_rays, ...) to (N_rays, N_samples, ...)
            cam_idxs = inputs['cam_idxs']
            cam_idxs = cam_idxs.reshape(N_rays, 1, -1).expand(-1, N_samples, -1)
            framecodes = self.framecodes(cam_idxs.reshape(N_rays * N_samples, -1))
            framecodes = framecodes.reshape(N_rays, N_samples, -1)
            view_inputs = torch.cat([view_inputs, framecodes], dim=-1)

        if 'normal' in encoded_pts and self.pred_sdf:
            normal = encoded_pts['normal']
            view_inputs = torch.cat([view_inputs, normal], dim=-1)

        view_inputs = view_inputs.flatten(end_dim=1)
        if 'valid_idxs' in encoded_pts:
            view_inputs = view_inputs[encoded_pts['valid_idxs']]

        return view_inputs, encoded

    def fill_prediction(self, preds: torch.Tensor, shape: Union[List, tuple], valid_info: Mapping[str, Any]):
        """ Create a full tensor from valid prediction.
        In evaluate_pts, we may avoid computation on some points that do not require prediction.
        The preds tensor shape is thus varying. This function turn the prediction into fixed size
        so that the prediction can be processed with batch operation.

        Args:
            preds (torch.Tensor): prediction tensor
            shape (Union[List, tuple]): shape of the first dimension in preds
            valid_info (Mapping[str, Any]): information that requires to map preds back to a tensor of shape [*shape, pred_size]

        Returns:
            outputs (torch.Tensor): tensor of shape [*shape, pred_size]
        """
        output_dim = preds.shape[-1]

        if valid_info is not None and 'valid_idxs' in valid_info:
            valid_idxs = valid_info['valid_idxs']
            outputs = torch.zeros(np.prod(shape), output_dim)
            # by default force nothing there in the empty space

            # outputs[invalid_idxs, -1] = 0.
            outputs[valid_idxs] += preds
            outputs = outputs.reshape(*shape, output_dim)
        else:
            outputs = preds.reshape(*shape, output_dim)

        return outputs

    def inference_sigma(self, density_inputs: torch.Tensor):
        """ Inference the geometry representation sigma (density or sdf) values
        """
        h = self.forward_density(density_inputs)
        sigma = self.sigma_linear(h)
        return sigma, h

    def inference_rgb(self, view_inputs: torch.Tensor, density_feature: torch.Tensor, rgb_eps: float = 0.001, **kwargs):
        """ Inference the color for the given input points
        """
        rgb = self.forward_view(view_inputs, density_feature)
        rgb = torch.sigmoid(rgb) * (1 + 2 * rgb_eps) - rgb_eps
        return rgb

    def forward_density(self, density_inputs: torch.Tensor):
        h = density_inputs

        for i, l in enumerate(self.pts_linears):
            h = self.pts_linears[i](h)
            h = F.relu(h, inplace=True)
            if i in self.skips:
                h = torch.cat([density_inputs, h], -1)
        return h

    def forward_view(self, view_inputs: torch.Tensor, density_feature: torch.Tensor):
        # produce features for color/radiance
        feature = self.feature_linear(density_feature)
        h = torch.cat([feature, view_inputs], -1)

        for i, l in enumerate(self.views_linears):
            h = self.views_linears[i](h)
            h = F.relu(h, inplace=True)

        return self.rgb_linear(h)

    def forward_geometry(
        self,
        render_data: Mapping[str, Any],
        chunk: int = 1024*64*5,
        return_gradient: bool = False,
        geom_extra_rets: List[str] = [],
        **kwargs
    ):
        """ Forward function for geometry only.

        Args:
            render_data (Mapping[str, Any]): dictionary of rendering data
            chunk (int): chunk size for forward to avoid OOM
            return_gradient (bool): whether to return gradient
            geom_extra_rets (List[str]): list of keys of extra tensors to return

        Returns:
            Dictionary of outputs.
        """

        assert 'pts' in render_data, f'Query locations need to be provided for forward_geometry'
        if 'kp3d' not in render_data:
            assert self.skel_type == SMPLSkeleton, f'Only SMPLSkeleton is supported for kinematic forward!'
            kp3d, skts = calculate_kinematic(
                self.rest_pose,
                render_data['bones'],
                render_data.get('root_locs', None),
            )
            render_data['skts'] = skts
            render_data['kp3d'] = kp3d

        pts = render_data['pts']
        kp3d = render_data['kp3d']
        skts = render_data['skts']
        bones = render_data['bones']

        # only deals with one pose at a time
        assert len(kp3d) == 1
        assert len(skts) == 1
        assert len(bones) == 1

        pts_shape = pts.shape
        pts = pts.reshape(1, -1, 3)
        N_samples = pts.shape[1]

        # pre-allocate
        density = torch.zeros(N_samples, 1)
        sigma = torch.zeros(N_samples, 1)
        extra_dict = {k: None for k in geom_extra_rets}
        if self.pred_sdf:
            sigma += 10.
        gradient = None
        if return_gradient:
            gradient = torch.zeros(N_samples, 3)
        for i in range(0, N_samples, chunk):
            chunk_pts = pts[:, i:i+chunk]
            geom_inputs = {
                'kp3d': kp3d,
                'skts': skts,
                'bones': bones,
                'N_unique': 1,
                'pts': chunk_pts,
            }
            _, preds = self.evaluate_pts(
                geom_inputs,
                geometry_only=True,
                geom_extra_rets=geom_extra_rets,
            )
            if preds is None:
                continue
            preds = {k: v.detach() for k, v in preds.items()}
            if 'valid_idxs' in preds:
                density[i + preds['valid_idxs']] = preds['density']
                sigma[i + preds['valid_idxs']] = preds['sigma']
                if gradient is not None:

                    pred_gradient = preds['surface_gradient'].flatten(end_dim=1)
                    pred_gradient = pred_gradient[preds['valid_idxs']]
                    gradient[i + preds['valid_idxs']] = pred_gradient

                for k in geom_extra_rets:
                    extra_tensor = preds[k]
                    if extra_dict[k] is None:
                        extra_dict[k] = torch.zeros(N_samples, extra_tensor.shape[-1])
                    extra_dict[k][i + preds['valid_idxs']] = extra_tensor
            else:
                density[i:i+chunk] = preds['density']
                sigma[i:i+chunk] = preds['sigma']
                if gradient is not None:
                    gradient[i:i+chunk] = preds['surface_gradient'].flatten(end_dim=1)

                for k in geom_extra_rets:
                    extra_tensor = preds[k]
                    if extra_dict[k] is None:
                        extra_dict[k] = torch.zeros(N_samples, extra_tensor.shape[-1])
                    extra_dict[k][i:i+chunk] = extra_tensor

        outputs = {
            'density': density.reshape(*pts_shape[:-1], 1),
            'sigma': sigma.reshape(*pts_shape[:-1], 1),
            **extra_dict,
        }

        if gradient is not None:
            outputs.update(gradient=gradient.reshape(*pts_shape[:-1], 3))
        return outputs

    def collect_encoded(self, encoded_pts: Mapping[str, torch.Tensor], encoded_views: Mapping[str, torch.Tensor]):
        """ Collect encodings that are useful for loss computations or rendering
        """
        ret = {}
        if self.pred_sdf and 'surface_gradient' in encoded_pts:
            ret['surface_gradient'] = encoded_pts['surface_gradient']
        return ret

    def raw2outputs(
        self,
        raw: torch.Tensor,
        z_vals: torch.Tensor,
        rays_d: torch.Tensor,
        bgs: Optional[torch.Tensor] = None,
        **kwargs
    ):
        # TODO: update
        """Transforms model's predictions to semantically meaningful values.
        Args:
            raw: [num_rays, num_samples along ray, 4]. Prediction from model.
            z_vals: [num_rays, num_samples along ray]. Integration time.
            rays_d: [num_rays, 3]. Direction of each ray.
            bgs: [num_rays, 3]. Background color
            act_fn: activation function for the density
        Returns:
            rgb_map: [num_rays, 3]. Estimated RGB color of a ray.
            disp_map: [num_rays]. Disparity map. Inverse of depth map.
            acc_map: [num_rays]. Sum of weights along each ray.
            weights: [num_rays, num_samples]. Weights assigned to each sampled color.
            depth_map: [num_rays]. Estimated distance to object.
        """
        if self.pred_sdf:
            return self.raw2sdfoutputs(raw, z_vals, rays_d, bgs=bgs, **kwargs)

        #raw2alpha = lambda raw, dists, noise, act_fn=act_fn: 1.-torch.exp(-(act_fn(raw + noise))*dists)

        dists = z_vals[...,1:] - z_vals[...,:-1]
        dists = torch.cat([dists, torch.Tensor([1e10]).expand(dists[...,:1].shape)], -1)  # [N_rays, N_samples]

        # Multiply each distance by the norm of its corresponding direction ray
        # to convert to real world distance (accounts for non-unit directions).
        dists = dists * torch.norm(rays_d[...,None,:], dim=-1)

        rgb = raw[..., :3]

        alpha = 1. - torch.exp(-raw[..., 3] * dists)

        # weights = alpha * tf.math.cumprod(1.-alpha + 1e-10, -1, exclusive=True)
        # sum_{i=1 to N samples} prob_of_already_hit_particles * alpha_for_i * color_for_i
        # C(r) = sum [T_i * (1 - exp(-sigma_i * delta_i)) * c_i] = sum [T_i * alpha_i * c_i]
        # alpha_i = 1 - exp(-sigma_i * delta_i)
        # T_i = exp(sum_{j=1 to i-1} -sigma_j * delta_j) = torch.cumprod(1 - alpha_i)
        # standard NeRF
        weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]
        rgb_map = torch.sum(weights[...,None] * rgb, -2)  # [N_rays, 3]

        depth_map = torch.sum(weights * z_vals, -1)
        disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / (torch.sum(weights, -1)  + 1e-10))

        invalid_mask = torch.ones_like(disp_map)
        invalid_mask[torch.isclose(weights.sum(-1), torch.tensor(0.))] = 0.
        disp_map = disp_map * invalid_mask

        acc_map = torch.minimum(torch.sum(weights, -1), torch.tensor(1.))

        if bgs is not None:
            rgb_map = rgb_map + (1. - acc_map)[..., None] * bgs

        return {'rgb_map': rgb_map, 'disp_map': disp_map, 'acc_map': acc_map,
                'weights': weights, 'alpha': alpha}

    def raw2sdfoutputs(
        self,
        raw: torch.Tensor,
        z_vals: torch.Tensor,
        rays_d: torch.Tensor,
        bgs: Optional[torch.Tensor] = None,
        **kwargs
    ):
        """ Sligtly different from usual case
        """

        dists = z_vals[...,1:] - z_vals[...,:-1]
        dists = torch.cat([dists, torch.Tensor([1e10]).expand(dists[...,:1].shape)], -1)  # [N_rays, N_samples]

        # Multiply each distance by the norm of its corresponding direction ray
        # to convert to real world distance (accounts for non-unit directions).
        dists = dists * torch.norm(rays_d[...,None,:], dim=-1)

        # following Lior's paper
        density = raw[..., 3]
        # note: this is in logspace, originally we do exp(-density*dists)
        free_energy = density * dists
        shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1), free_energy[:, :-1]], dim=-1)  # shift one step

        # weights = alpha * tf.math.cumprod(1.-alpha + 1e-10, -1, exclusive=True)
        # sum_{i=1 to N samples} prob_of_already_hit_particles * alpha_for_i * color_for_i
        # C(r) = sum [T_i * (1 - exp(-sigma_i * delta_i)) * c_i] = sum [T_i * alpha_i * c_i]
        # alpha_i = 1 - exp(-sigma_i * delta_i)
        # T_i = exp(sum_{j=1 to i-1} -sigma_j * delta_j) = torch.cumprod(1 - alpha_i)
        # standard NeRF
        alpha = 1 - torch.exp(-free_energy)  # probability of it is not empty here
        T = torch.exp(-torch.cumsum(shifted_free_energy, dim=-1))  # probability of everything is empty up to now
        weights = alpha * T # probability of the ray hits something here
        #weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]

        # get rgb (as usual)
        rgb = raw[..., :3] # [N_rays, N_samples, 3]
        rgb_map = torch.sum(weights[...,None] * rgb, -2)  # [N_rays, 3]

        eps = 1e-10
        inv_eps = 1. / eps
        depth_map = torch.sum(weights * z_vals, -1)
        disp_map = (weights.sum(dim=-1) / depth_map).double()
        disp_map = torch.where(
            (disp_map > 0) & (disp_map < inv_eps) & (weights.sum(dim=-1) > eps), disp_map, inv_eps
        ).float()
        encoded = kwargs['encoded']
        """
        disp_map = 1./torch.max(eps * torch.ones_like(depth_map), depth_map / (torch.sum(weights, -1)  + eps))

        invalid_mask = torch.ones_like(disp_map)
        invalid_mask[torch.isclose(weights.sum(-1), torch.tensor(0.))] = 0.
        disp_map = disp_map * invalid_mask
        """

        acc_map = torch.minimum(torch.sum(weights, -1), torch.tensor(1.))

        if bgs is not None:
            rgb_map = rgb_map + (1. - acc_map)[..., None] * bgs

        return {'rgb_map': rgb_map, 'disp_map': disp_map, 'acc_map': acc_map,
                'weights': weights, 'alpha': alpha}

    def collect_outputs(
        self,
        ret: Mapping[str, torch.Tensor],
        ret0: Optional[Mapping[str, torch.Tensor]] = None,
        encoded: Optional[Mapping[str, torch.Tensor]] = None,
        encoded0: Optional[Mapping[str, torch.Tensor]] = None
    ):
        """ Collect outputs into a dictionary for loss computation/rendering

        Parameter
        ---------
        ret: dictionary of the fine-level rendering
        ret0: dictionary of the coarse-level rendering
        encoded: dictionary of the fine-level model info
        encoded0: dictionary of the corase-level model info

        """

        collected = {'rgb_map': ret['rgb_map'], 'disp_map': ret['disp_map'],
                     'acc_map': ret['acc_map'],}
        if not self.training:
            return collected

        collected['T_i'] = ret['weights']
        collected['alpha'] = ret['alpha']

        if ret0 is not None:
            collected['rgb0'] = ret0['rgb_map']
            collected['disp0'] = ret0['disp_map']
            collected['acc0'] = ret0['acc_map']
            collected['alpha0'] = ret0['alpha']

        if 'j_dists' in ret and self.training:
            collected['j_dists'] = ret['j_dists']
            if ret0 is not None:
                collected['j_dists0'] = ret0['j_dists']

        if encoded is not None and 'surface_gradient' in encoded and self.training:
            collected['surface_gradient'] = encoded['surface_gradient']

            if encoded0 is not None:
                collected['surface_gradient0'] = encoded0['surface_gradient']

        if self.training and 'bgs' in ret:
            collected['bgs'] = ret['bgs']
            collected['bg_preds'] = ret['bg_preds']

        return collected

    def _empty_outputs(self, batch: Mapping[str, Any]):
        N_rays = len(batch['rays_o'])
        empty = torch.zeros(N_rays, 3)
        rgb_empty = batch.get('bgs', empty)
        return {'rgb_map': rgb_empty, 'disp_map': empty[:, 0], 'acc_map': empty[:, 0]}

    def to_density(self, sigma: torch.Tensor):
        """ Turn sigma from raw prediction (value unbounded) to density (value >= 0.)
        """
        if self.pred_sdf:
            return self.sdf2density(sigma)

        noise = 0.
        if self.training and self.density_noise_std > 0.:
            noise = torch.randn_like(sigma) * self.density_noise_std

        sigma = sigma + noise

        return F.relu(sigma, inplace=True)

    def sdf2density(self, sdf: torch.Tensor):
        # use laplace CDF to represent density, with sdf being pdf
        B = (self.B.abs() + 0.0001)
        inv_B = 1 / B
        return inv_B * (0.5 + 0.5 * sdf.sign() * torch.expm1(-sdf.abs() / B))


inputs = []

for i in range(9):
    user_input = input("Input: ")

    if user_input not in ['x', 'o']:
        print("Wrong input")
        exit()

    inputs.append(user_input)

print("--------")
for i in range(3):
    print("|%s|%s|%s|" % (inputs[3*i], inputs[3*i+1], inputs[3*i+2]))
    print("--------")


import random
import torch
import torch.nn.functional as F

def beamSearch():
  None

def random_sampling(predict, vocab, k = 1024):
  # k개 중 랜덤으로 선택된 값을 반환. 사실상 top_k와 같은 기
  gen =[]

  probs, indexs = torch.topk(predict, k=k, dim=-1)
  probs = probs.squeeze().tolist()[-1]
  indexs = indexs.squeeze().tolist()[-1]

  for i in range(len(indexs)):
    gen.append((vocab.to_tokens(indexs[i]),probs[i]))
  # print('topk word and value: ', gen)

  rand_num = random.randint(0,k-1)
  gen_word = vocab.to_tokens(indexs[rand_num])

  return gen_word

def top_p(logits, vocab, threshold = 0.9):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    indexs = sorted_indices.tolist()

    sorted_softmax_logits = F.softmax(sorted_logits, dim=-1)
    cum_probs = torch.cumsum( sorted_softmax_logits, dim=-1)

    sorted_indices_to_remove = cum_probs > threshold
    top_p_index = 0

    # Top-p에 해당하는 index를 획득
    for i in range(len(sorted_indices_to_remove)):
      if sorted_indices_to_remove[i]== True:
        top_p_index = 0 if i==0 else i-1
        break

    # for i in range(top_p_index):
      # print('gen '+str(i)+': '+vocab.to_tokens(indexs[i]))

    rand_num = random.randint(0, top_p_index) # top-p 분포에서 랜덤 샘플링
    top_p_sample_num = indexs[rand_num]
    gen_word = vocab.to_tokens(top_p_sample_num)
    print('selected token: '+gen_word+ ' softmax value:'+str(sorted_softmax_logits[rand_num]))

    return gen_word

def top_k(predict, vocab, k):
  # topk 중 랜덤으로 선택된 값을 반환.
  gen = []

  probs, indexs = torch.topk(predict, k=k,dim=-1)
  # probs = probs.squeeze().tolist()[-1]
  # indexs = indexs.squeeze().tolist()[-1]
  probs = probs.tolist()
  indexs = indexs.tolist()

  for i in range(len(indexs)):
    gen.append((vocab.to_tokens(indexs[i]), probs[i]))
  # print('topk word and value: ', gen)

  rand_num = random.randint(0, k - 1)
  gen_word = vocab.to_tokens(indexs[rand_num])

  return gen_word


import string
#open txt
#Open Tianyi Yang's text
with open('Tianyi.txt') as fp:
    data1 = fp.readlines()
#Open Mengge Geng's text
with open('Mengge.txt') as fp:
    data2 = fp.readlines()
#Open Jiefu Dong's text
with open('Jiefu.txt') as fp:
    data3 = fp.readlines()
#Open Han Luo's text
with open('Han.txt') as fp:
    data4 = fp.readlines()
allData = data1+data2+data3+data4 #Merging all files together

def convert(myString): #Remove punctuations
    mylist = myString
    for i in range (0,len(mylist)):
        mylist[i]=mylist[i].replace("\n", "")
        for c in mylist[i]:
            if c in string.punctuation:
                mylist[i] = mylist[i].replace(c,"")
    return(mylist)

newData = str(convert(allData))

if __name__ == '__main__':
    with open('Compile.txt', 'w') as fp:
        fp.write(newData)


#!/usr/bin/env python3

import argparse

import matplotlib.pyplot as plt

from sipmarray import SiPMarray

def parse_args():
    parser = argparse.ArgumentParser(
        description=('Script to quickly get a SiPM array plot or geometry.'))

    parser.add_argument('-m', '--model',
                        help='Define the SiPM model.',
                        default= '6x6',
                        type = str,
                        required=True)
    parser.add_argument('-d', '--diameter',
                        help='Define the SiPM array diameter in mm.',
                        required= True,
                        type= float)
    parser.add_argument('-b', '--border',
                        help='Define the SiPM array border margin in mm.',
                        required= False,
                        type= float,
                        default= 0)
    parser.add_argument('-p', '--plot',
                        help='Plot the SiPM array.',
                        type= bool,
                        default= False,
                        required=False)
    parser.add_argument('-o', '--output',
                        help ='Output the SiPM array geometry to a file.',
                        type = bool,
                        default = True,
                        required=False)
    parser.add_argument('-v', '--verbose',
                        help ='Print the SiPM array properties in the terminal.',
                        type = bool,
                        default = False,
                        required=False)
    args = parser.parse_args()

    model = args.model
    diameter = args.diameter
    border = args.border
    plot = args.plot
    output = args.output
    verbose = args.verbose
    return model,diameter,border,plot,output,verbose

def make_plot(model, diameter, array):
    fig, ax = plt.subplots(1,1)
    fig, ax = array.plot_sipm_array(figax= (fig,ax))
    fig.savefig(f'sipmarray_{model}_{diameter}mm.png')

def make_output(model, diameter, array):
    array.export_corners_active(
            file_name = f'sipmarray_corners_active_{model}_{diameter}mm.csv')
    array.export_corners_package(
            file_name = f'sipmarray_corners_package_{model}_{diameter}mm.csv')
    array.export_centres(
            file_name = f'sipmarray_centres_{model}_{diameter}mm.csv')

def main():
    model, diameter, border, plot, output, verbose = parse_args()
    if verbose: print('Building SiPM array.')

    array = SiPMarray(array_diameter = diameter,
                      border_margin = border,
                      sipm_model = model)

    if verbose: array.print_properties(unit_properties=True)
    if plot: make_plot(model, diameter, array)
    if output: make_output(model, diameter, array)
    if verbose: print('Done!')

if __name__ == '__main__':
    main()


from typing import Optional

import numpy as np
import scipy.sparse
from scipy.sparse import csr_matrix
from sqlalchemy import text

class Popularity:
    def __init__(self, window_size: int = 10, retrain_interval: int = 10):
        self.window_size = window_size
        self.retrain_interval = retrain_interval
        # Popularity matrix
        self.items_popularity: Optional[scipy.sparse.csr_matrix] = None

    def run(self, x: scipy.sparse.csr_matrix) -> Optional[csr_matrix]:
        return self.__fit(x).__predict(x)

    def retrain(self, date, dataset, k):
        # Return a query that'll return the Top-K most bought items of the past {window_size} days
        return text(
            f'SELECT item_id FROM (\n'
            f'SELECT item_id,\n' +
            f'       sum(item_count) OVER (PARTITION BY item_id ORDER BY purchase_date\n' +
            f'           ROWS BETWEEN {self.window_size} PRECEDING AND CURRENT ROW)\n' +
            f'           as total_interactions\n' +
            f'FROM items_day_count\n' +
            f'WHERE purchase_date = \'{date}\' AND dataset_name=\'{dataset}\'\n'
            f'ORDER BY total_interactions DESC LIMIT {k}) as tmp;'
        )

    def __fit(self, x: scipy.sparse.csr_matrix) -> 'Popularity':
        # Get popularity of each item
        items_popularity = np.asarray(x.sum(0))
        # Create index for each item
        col_index: np.ndarray = np.arange(items_popularity.shape[1])
        # Add index row as 2'nd row to array
        items_popularity: np.ndarray = np.vstack([items_popularity, col_index])
        # Sort array by values in first row in descending order
        result: np.ndarray = items_popularity[:, np.flip(np.argsort(items_popularity[0]))]
        # Applying window size
        result = result[:, :self.window_size]
        # result = np.resize(result, (result.shape[0], self.window_size))  # Cutting all products that aren't wanted
        self.items_popularity = result
        return self

    def __predict(self, histories: scipy.sparse.csr_matrix) -> Optional[csr_matrix]:
        return self.items_popularity

if __name__ == "__main__":
    a = Popularity(window_size=5)
    print(a.run(scipy.sparse.csr_matrix([[10, 2, 12], [1, 2, 3], [1, 2, 3], [5, 1, 8], [1, 9, 4], [9, 9, 9]])))


# -*- coding: utf-8 -*-
import cvxpy as cp
import numpy as np
from scipy import stats
from sympy.utilities.iterables import multiset_permutations
from scipy.special import comb
from scipy.optimize import fmin_l_bfgs_b

class BinLogitCMLE():
    """
    Implement the conditional maximum likelihood estimator of the binary logit
    model with fixed effects (T balanced periods).

    Reference : Woolridge, "Econometric Analysis of Cross-Section and Panel Data" (Chap 15.8.3).

    Inputs
    ------------
    `A': (n x T x K) numpy.array that contains individuals' covariates at each time period;
    `b': (n x T) numpy.array that contains individuals' choices at each time period.
    """

    def __init__(self, A, b):
        self.A = A
        # make sure b is of correct format
        # make sure it contains only ones
        # make sure A is numeric
        self.b = b
        if len(A.shape) != 3:
            print('The design matrix A must have shape (N x T x K).')
        else:
            self.n, self.T, self.K = A.shape

        self.R = self.compute_perm()
        if self.n<= 100:
            existence_cond = self.separation_test(verbose=False)
            if existence_cond==False:
                print("Warning ! Data violate CMLE's existence conditions.")

    def compute_perm(self):
        """
        Returns
        `R' : numpy.ndarray that contains all vectors of 1/0 of size T such
        that their coordinates sum equals k, for k in [0,T]
        """
        R = list()
        for k in range(self.T + 1):
            array = np.ndarray(shape=(int(comb(self.T, k)), self.T))
            select = np.zeros(self.T)
            select[:k] = np.ones(k)
            for idx, p in enumerate(multiset_permutations(select)):
                array[idx] = p
            R.append(array)
        return R

    def objective_i(self, i, A, b, beta):
        """
        Returns the i-th contribution to the objective function
        (CMLE is a finite-sum).

        Inputs
        ------------
        `A' : design matrix;
        `b' : outcome variable;
        """
        R = self.R
        Xprime_beta = A[i].dot(beta)
        n_i = int(np.sum(b[i]))
        omega_i = np.sum(np.exp(R[n_i].dot(Xprime_beta)), axis=0)
        res = np.dot(b[i], Xprime_beta) - np.log(omega_i)
        return res

    def objective(self, beta, A, b):
        objective = 0
        for i in range(self.n):
            objective += self.objective_i(i, A, b, beta)
        return objective / self.n

    def loss(self, beta, A, b):
        objective = 0
        for i in range(self.n):
            objective += self.objective_i(i, A, b, beta)
        return - objective / self.n

    def loss_grad(self, beta, A, b):
        """
        Computes the global gradient.
        """
        g = np.zeros_like(beta)
        for i in range(self.n):
            g += self.comp_grad_i(i, beta)
        return - g / self.n

    def comp_grad_i(self, i, beta):
        """
        Returns the pointwise gradient evaluated on datum i.
        """
        A = self.A
        b = self.b
        R = self.R
        n_i = int(np.sum(b[i]))
        Xprime_beta = A[i].dot(beta)
        if ((n_i!=0) & (n_i!=self.T)):
            omega_i = (1. / (np.sum(np.exp(R[n_i].dot(Xprime_beta)))))
            g = (np.dot(b[i], A[i]) - np.sum(np.dot(R[n_i], A[i]) *
                        np.tile(np.exp(R[n_i].dot(Xprime_beta)),
                                (self.K, 1)).T, axis=0) * omega_i)
        else:
            g = np.zeros(self.K)
        return g

    def comp_hessian_i(self, i, beta):
        """
        Returns the pointwise Hessian matrix evaluated on datum i.
        """
        A = self.A
        b = self.b
        R = self.R
        n_i = int(np.sum(b[i]))
        Xprime_beta = A[i].dot(beta)
        if ((n_i!=0) & (n_i!=self.T)):
            omega_i = (1. / (np.sum(np.exp(R[n_i].dot(Xprime_beta)))))
            b_i = np.sum(np.dot(R[n_i], A[i]) *
                        np.tile(np.exp(R[n_i].dot(Xprime_beta)), (self.K, 1)).T, axis=0)
            hess = (np.outer(b_i, b_i) * omega_i ** 2 -
                    np.sum([np.outer(elem, elem) * np.exp(R[n_i].dot(Xprime_beta))[idx] for idx, elem in enumerate(np.dot(R[n_i], A[i]))])
                    * omega_i)
        else:
           hess= np.zeros((self.K, self.K))
        return hess

    def comp_gradient(self, beta, A, b):
        """
        Computes the global gradient.
        """
        g = np.zeros_like(beta)
        for i in range(self.n):
            g += self.comp_grad_i(i, beta)
        return g / self.n

    def comp_hessian(self, beta): # warning : one-dimensional case
        """Computes the global Hessian.
        """
        hess = np.ndarray(shape=(self.K,self.K), buffer=np.zeros(self.K**2))
        for i in range(self.n):
            hess = hess + self.comp_hessian_i(i, beta)
        return hess / self.n

    def fit(self, beta_init, n_iter=1000, step=0.01, tol=1e-6, epsilon=1e-10, hessian=False, verbose=True, BFGS=True):
        """Fit the conditional logit using either BFGS algorithm or
        the Newton-Raphson algorithm (with or without an Hessian step).
        """
        beta = beta_init.copy()
        beta_list = [beta_init]
        if BFGS:
            beta, f_min, _ = fmin_l_bfgs_b(self.loss, beta_init, self.loss_grad, args=(self.A, self.b), pgtol=1e-6, factr=1e-30)
        elif hessian:
            while True:
                beta_t = beta - np.linalg.inv(self.comp_hessian(beta)).dot(self.comp_gradient(beta, self.A, self.b))
                t = abs(beta_t - beta)
                if t < tol:
                    break
                beta = beta_t
                #beta_list.append(beta)
        else:
            stop = True
            for t in range(n_iter):
                if stop:
                    beta = beta + step * self.comp_gradient(beta, self.A, self.b)
                    if verbose:
                        print("Iteration %s completed" % t)
                    beta_list.append(beta)
                    if self.objective(beta, self.A, self.b, ) -self.objective( beta_list[t], self.A, self.b,) < epsilon:
                        stop = False
        return beta, beta_list

    def AsympVariance(self, beta, score=True):
        avar = np.ndarray(shape=(self.K,self.K), buffer=np.zeros(self.K**2))
        if score:
            for i in range(self.n):
                grad = self.comp_grad_i(i, beta)
                avar += np.outer(grad, grad)
        else:
            for i in range(self.n):
                avar += - self.comp_hessian_i(i, beta) # Fix the minus ? -> get negative variance
        return np.linalg.inv(avar / self.n)

    def lr_nulltest(self, bet_hat, lvl=0.05):
         """
         Perform the LR test of the global null.
         ----------
         Outputs : LR statistics, p-value
         """
         lu = self.objective(bet_hat, self.A, self.b)
         lr = self.objective(np.zeros(self.K), self.A, self.b)
         ratio = 2 *self.n * (lu - lr)
         pval = stats.chi2.sf(ratio, self.K)
         if pval<lvl:
             print("H0 : \beta_j =0 \forall j is rejected at level : ", lvl)
         return ratio, pval

    def separation_test(self, tol=1e-4, verbose=True):
        """
        Verify that CMLE's existence condition is met (False=CMLE does not exist).
        """
        res = True
        R = self.R
        rn = np.sum([len(R[int(np.sum(self.b[i]))]) for i in range(self.n)])
        u = cp.Variable(rn)
        a = []
        for i in range(self.n):
            for j in range(len(R[int(np.sum(self.b[i]))])):
                vec = R[int(np.sum(self.b[i]))][j] - self.b[i]
                a.append(np.sum(vec[:,None] * self.A[i]))
        a = np.array(a)
        prob = cp.Problem(cp.Minimize(cp.norm(cp.sum(cp.multiply(u,a),
                                                axis=0))), [u>=np.ones(rn)])
        prob.solve()
        if np.abs(prob.value)>tol:
            res = False
        if verbose:
            print('CMLE does exist: {}.'.format(res))
        return res


from .ICorrect import ICorrect
from modules.unscented_utils import *

class UnscentedCorrect(ICorrect):

    def __init__(self, f_func, h_func, f_args, h_args,  Wm, Wc, lambda_,  d=2, image_size= (4094, 2046)):

        self.d = d
        self.lambda_ = lambda_
        self.h_func = h_func
        self.f_func = f_func
        self.Wm = Wm
        self.Wc = Wc
        self.image_size = image_size
        self.f_args =f_args
        self.h_args =h_args

    def define_params(self, Xs):
        self.Xs = Xs

    def correct(self, z, R, pred_state, pred_cov, state, state_cov, delta_t):
        Ys = sigma_points(pred_state, pred_cov, lambda_=self.lambda_, N=self.d)
        pred_z = propagate_func_corr(self.h_func, self.Wm, self.Wc, Ys, delta_t,
                                          args=self.h_args,image_size=self.image_size) # z^, S_k - R

        residuals = np.zeros(shape=state.shape)
        residuals[0] = z - pred_z[0]
        residuals[1] = pred_z[1]

        #Innovation
        S_innovation = propagate_func_corr(self.h_func, self.Wm, self.Wc, Ys, delta_t,
                                          args=self.h_args,image_size=self.image_size,
                                          mean=False)
        S_innovation[0] = S_innovation[0] + R

        # Cross covariance matrix
        C = cross_covariance(self.f_func, self.h_func, self.Wc, self.Xs, Ys, delta_t, self.f_args,
                         self.h_args, pred_state, pred_z, image_size=self.image_size)

        K = optimal_gain(C, S_innovation, image_size=self.image_size)

        k = np.zeros((tuple([2])+self.image_size))
        k[0] = K[0]*residuals[0] + K[1]*residuals[1]
        k[1] = K[2]*residuals[0] + K[3]*residuals[1]
        state = pred_state + k

        KSKt = get_KSKt_product(K, S_innovation, image_size=self.image_size)
        state_cov = pred_cov - KSKt

        return state, state_cov, K[[0, 3]]


import re
from collections import OrderedDict
from typing import Union, Optional, Any, Dict, List, Type, TYPE_CHECKING
from .base import DEFINES, ENUMS, TYPEDEFS, STRUCTS
from .field import calculate_padding, Kind, FieldType
from .c_expr import c_eval
from .exceptions import CStructException, ParserError
from .native_types import get_native_type

if TYPE_CHECKING:
    from .abstract import AbstractCStruct, AbstractCEnum

__all__ = ["parse_struct", "parse_struct_def", "parse_enum_def", "Tokens"]

SEPARATORS = [" ", "\t", "\n", ";", "{", "}", ":", ",", "="]
SPACES = [" ", "\t", "\n"]

class Tokens(object):
    def __init__(self, text: str) -> None:
        # remove the comments
        text = re.sub(r"//.*?$|/\*.*?\*/", "", text, flags=re.S | re.MULTILINE)
        # c preprocessor
        lines = []
        for line in text.split("\n"):
            if re.match(r"^\s*#define", line):
                try:
                    _, name, value = line.strip().split(maxsplit=2)
                    DEFINES[name] = c_eval(value)
                except Exception:
                    raise ParserError(f"Parsing line `{line}`")
            else:
                lines.append(line)
        text = "\n".join(lines)
        self.tokens = self.tokenize(text)

    def tokenize(self, text) -> List[str]:
        tokens: List[str] = []
        t: List[str] = []
        for c in text:
            if c in SEPARATORS:
                if t:
                    tokens.append("".join(t))
                    t.clear()
                if c not in SPACES:
                    tokens.append(c)
            else:
                t.append(c)
        if t:
            tokens.append(t.getvalue())
        return tokens

    def pop(self) -> str:
        return self.tokens.pop(0)

    def pop_c_type(self) -> str:
        c_type = self.pop()
        if c_type in ["signed", "unsigned"] and len(self) > 1:
            # short int, long int, or long long
            c_type = c_type + " " + self.pop()
        elif c_type in ["short", "long"] and len(self) > 1 and self.get() in ["int", "long"]:
            # short int, long int, or long long
            c_type = c_type + " " + self.pop()
        return c_type

    def get(self) -> str:
        return self.tokens[0]

    def push(self, value: str) -> None:
        return self.tokens.insert(0, value)

    def __len__(self) -> int:
        return len(self.tokens)

    def __str__(self) -> str:
        return str(self.tokens)

def parse_type(tokens: Tokens, __cls__: Type["AbstractCStruct"], byte_order: Optional[str], offset: int) -> "FieldType":
    if len(tokens) < 2:
        raise ParserError("Parsing error")
    c_type = tokens.pop()
    # signed/unsigned/struct
    if c_type in ["signed", "unsigned", "struct", "union", "enum"] and len(tokens) > 1:
        c_type = c_type + " " + tokens.pop()

    vlen = 1
    flexible_array = False

    if not c_type.endswith("{"):
        next_token = tokens.pop()
        # short int, long int, or long long
        if next_token in ["int", "long"]:
            c_type = c_type + " " + next_token
            next_token = tokens.pop()
        # void *
        if next_token.startswith("*"):
            next_token = next_token[1:]
            c_type = "void *"
        # parse length
        if "[" in next_token:
            t = next_token.split("[")
            if len(t) != 2:
                raise ParserError(f"Error parsing: `{next_token}`")
            next_token = t[0].strip()
            vlen_part = t[1]
            vlen_expr = []
            while not vlen_part.endswith("]"):
                vlen_expr.append(vlen_part.split("]")[0].strip())
                vlen_part = tokens.pop()
            t_vlen = vlen_part.split("]")[0].strip()
            vlen_expr.append(vlen_part.split("]")[0].strip())
            t_vlen = " ".join(vlen_expr)
            if not t_vlen:
                flexible_array = True
                vlen = 0
            else:
                try:
                    vlen = c_eval(t_vlen)
                except (ValueError, TypeError):
                    vlen = int(t_vlen)
        tokens.push(next_token)
        # resolve typedefs
        while c_type in TYPEDEFS:
            c_type = TYPEDEFS[c_type]

    # calculate fmt
    if c_type.startswith("struct ") or c_type.startswith("union "):  # struct/union
        c_type, tail = c_type.split(" ", 1)
        kind = Kind.STRUCT if c_type == "struct" else Kind.UNION
        if tokens.get() == "{":  # Named nested struct
            tokens.push(tail)
            tokens.push(c_type)
            ref: Union[Type[AbstractCEnum], Type[AbstractCStruct]] = __cls__.parse(tokens, __name__=tail, __byte_order__=byte_order)
        elif tail == "{":  # Unnamed nested struct
            tokens.push(tail)
            tokens.push(c_type)
            ref = __cls__.parse(tokens, __byte_order__=byte_order)
        else:
            try:
                ref = STRUCTS[tail]
            except KeyError:
                raise ParserError(f"Unknown `{c_type} {tail}`")
    elif c_type.startswith("enum"):
        from .cenum import CEnum

        c_type, tail = c_type.split(" ", 1)
        kind = Kind.ENUM
        if tokens.get() == "{":  # Named nested struct
            tokens.push(tail)
            tokens.push(c_type)
            ref = CEnum.parse(tokens, __name__=tail)
        elif tail == "{":  # unnamed nested struct
            tokens.push(tail)
            tokens.push(c_type)
            ref = CEnum.parse(tokens)
        else:
            try:
                ref = ENUMS[tail]
            except KeyError:
                raise ParserError(f"Unknown `{c_type} {tail}`")
    else:  # other types
        kind = Kind.NATIVE
        ref = None
    return FieldType(kind, c_type, ref, vlen, flexible_array, byte_order, offset)

def parse_typedef(tokens: Tokens, __cls__: Type["AbstractCStruct"], byte_order: Optional[str]) -> None:
    field_type = parse_type(tokens, __cls__, byte_order, 0)
    vname = tokens.pop()
    if field_type.ref is None:
        TYPEDEFS[vname] = field_type.c_type
    elif field_type.ref.__is_enum__:
        TYPEDEFS[vname] = f"enum {field_type.ref.__name__}"
    elif field_type.ref.__is_union__:
        TYPEDEFS[vname] = f"union {field_type.ref.__name__}"
    else:
        TYPEDEFS[vname] = f"struct {field_type.ref.__name__}"
    t = tokens.pop()
    if t != ";":
        raise ParserError(f"`;` expected but `{t}` found")

def parse_struct_def(
    __def__: Union[str, Tokens],
    __cls__: Type["AbstractCStruct"],
    __byte_order__: Optional[str] = None,
    process_muliple_definition: bool = False,
    **kargs: Any,  # Type['AbstractCStruct'],
) -> Optional[Dict[str, Any]]:
    # naive C struct parsing
    if isinstance(__def__, Tokens):
        tokens = __def__
    else:
        tokens = Tokens(__def__)
    result = None
    while tokens and (process_muliple_definition or not result):
        kind = tokens.pop()
        if kind == ";":
            pass

        elif kind == "typedef":
            if result:
                result["__cls__"].parse(result, **kargs)
            parse_typedef(tokens, __cls__, __byte_order__)

        elif kind == "enum":
            if result:
                result["__cls__"].parse(result, **kargs)
            name = tokens.pop()
            native_format = None
            if tokens.get() == ":":  # enumeration type declaration
                tokens.pop()  # pop ":"
                type_ = get_native_type(tokens.pop_c_type())
                native_format = type_.native_format
            if tokens.get() == "{":  # named enum
                tokens.pop()  # pop "{"
                result = parse_enum(tokens, __name__=name, native_format=native_format)
            elif name == "{":  # unnamed enum
                result = parse_enum(tokens, native_format=native_format)
            else:
                raise ParserError(f"`{name}` definition expected")

        elif kind in ["struct", "union"]:
            if result:
                result["__cls__"].parse(result, **kargs)
            __is_union__ = kind == "union"
            name = tokens.pop()
            if name == "{":  # unnamed nested struct
                result = parse_struct(tokens, __cls__=__cls__, __is_union__=__is_union__, __byte_order__=__byte_order__)
            elif tokens.get() == "{":  # Named nested struct
                tokens.pop()  # pop "{"
                result = parse_struct(
                    tokens, __cls__=__cls__, __is_union__=__is_union__, __byte_order__=__byte_order__, __name__=name
                )
            else:
                raise ParserError(f"`{name}` definition expected")

        else:
            raise ParserError(f"struct, union, or enum expected - `{kind}` found")
    return result

def parse_enum_def(__def__: Union[str, Tokens], **kargs: Any) -> Optional[Dict[str, Any]]:
    # naive C enum parsing
    if isinstance(__def__, Tokens):
        tokens = __def__
    else:
        tokens = Tokens(__def__)
    if not tokens:
        return None
    kind = tokens.pop()
    if kind not in ["enum"]:
        raise ParserError(f"enum expected - `{kind}` found")

    name = tokens.pop()
    native_format = None
    if tokens.get() == ":":  # enumeration type declaration
        tokens.pop()  # pop ":"
        type_ = get_native_type(tokens.pop_c_type())
        native_format = type_.native_format
    if tokens.get() == "{":  # named enum
        tokens.pop()  # pop "{"
        return parse_enum(tokens, __name__=name, native_format=native_format)
    elif name == "{":  # unnamed enum
        return parse_enum(tokens)
    else:
        raise ParserError(f"`{name}` definition expected")

def parse_enum(
    __enum__: Union[str, Tokens],
    __name__: Optional[str] = None,
    native_format: Optional[str] = None,
    **kargs: Any,
) -> Optional[Dict[str, Any]]:
    """
    Parser for C-like enum syntax.

    Args:
        __enum__:       definition of the enum in C syntax
        __name__:       enum name
        native_format:  struct module format

    Returns:
        dict: the parsed definition
    """
    from .cenum import CEnum

    constants: Dict[str, int] = OrderedDict()

    if isinstance(__enum__, Tokens):
        tokens = __enum__
    else:
        tokens = Tokens(__enum__)

    while len(tokens):
        if tokens.get() == "}":
            tokens.pop()
            break

        name = tokens.pop()
        next_token = tokens.pop()
        if next_token in {",", "}"}:  # enum-constant without explicit value
            if len(constants) == 0:
                value = 0
            else:
                value = next(reversed(constants.values())) + 1
        elif next_token == "=":  # enum-constant with explicit value
            exp_elems = []
            next_token = tokens.pop()
            while next_token not in {",", "}"}:
                exp_elems.append(next_token)
                if len(tokens) > 0:
                    next_token = tokens.pop()
                else:
                    break

            if len(exp_elems) == 0:
                raise ParserError("enum is missing value expression")

            int_expr = " ".join(exp_elems)
            try:
                value = c_eval(int_expr)
            except (ValueError, TypeError):
                value = int(int_expr)
        else:
            raise ParserError(f"`{__enum__}` is not a valid enum expression")

        if name in constants:
            raise ParserError(f"duplicate enum name `{name}`")
        constants[name] = value

        if next_token == "}":
            break

    result = {
        "__constants__": constants,
        "__is_struct__": False,
        "__is_union__": False,
        "__is_enum__": True,
        "__name__": __name__,
        "__native_format__": native_format,
        "__cls__": CEnum,
    }
    return result

def parse_struct(
    __struct__: Union[str, Tokens],
    __cls__: Type["AbstractCStruct"],
    __is_union__: bool = False,
    __byte_order__: Optional[str] = None,
    __name__: Optional[str] = None,
    **kargs: Any,
) -> Dict[str, Any]:
    """
    Parser for C-like struct syntax.

    Args:
        __struct__:     definition of the struct/union in C syntax
        __cls__:        base class (MemCStruct or CStruct)
        __is_union__:   True for union, False for struct
        __byte_order__: byte order, valid values are LITTLE_ENDIAN, BIG_ENDIAN, NATIVE_ORDER
        __name__:       struct/union name

    Returns:
        dict: the parsed definition
    """
    # naive C struct parsing
    from .abstract import AbstractCStruct
    from .mem_cstruct import MemCStruct

    if __cls__ is None or __cls__ == AbstractCStruct:
        __cls__ = MemCStruct
    __is_union__ = bool(__is_union__)
    fields_types: Dict[str, FieldType] = OrderedDict()
    flexible_array: bool = False
    offset: int = 0
    max_alignment: int = 0
    anonymous: int = 0
    if isinstance(__struct__, Tokens):
        tokens = __struct__
    else:
        tokens = Tokens(__struct__)
    while len(tokens):
        if tokens.get() == "}":
            tokens.pop()
            break
        # flexible array member must be the last member of such a struct
        if flexible_array:
            raise CStructException("Flexible array member must be the last member of such a struct")
        field_type = parse_type(tokens, __cls__, __byte_order__, offset)
        vname = tokens.pop()
        if vname in fields_types:
            raise ParserError(f"Duplicate member `{vname}`")
        if vname in dir(__cls__):
            raise ParserError(f"Invalid reserved member name `{vname}`")
        # anonymous nested union
        if vname == ";" and field_type.ref is not None and (__is_union__ or field_type.ref.__is_union__):
            # add the anonymous struct fields to the parent
            for nested_field_name, nested_field_type in field_type.ref.__fields_types__.items():
                if nested_field_name in fields_types:
                    raise ParserError(f"Duplicate member `{nested_field_name}`")
                # set the corret offset
                nested_field_type = nested_field_type.copy()
                nested_field_type.base_offset = offset + nested_field_type.base_offset
                nested_field_type.offset = offset + nested_field_type.offset
                fields_types[nested_field_name] = nested_field_type
            vname = f"__anonymous{anonymous}"
            anonymous += 1
            tokens.push(";")
        fields_types[vname] = field_type
        # calculate the max field size (for the alignment)
        max_alignment = max(max_alignment, field_type.alignment)
        # align struct if byte order is native
        if not __is_union__:  # C struct
            field_type.align_filed_offset()
            offset = field_type.offset + field_type.vsize
        t = tokens.pop()
        if t != ";":
            raise ParserError(f"`;` expected but `{t}` found")

    if __is_union__:  # C union
        # Calculate the sizeof union as size of its largest element
        size = max([x.vsize for x in fields_types.values()])
    else:  # C struct
        # add padding to struct if byte order is native
        size = offset + calculate_padding(__byte_order__, max_alignment, offset)

    # Prepare the result
    result = {
        "__fields__": list(fields_types.keys()),
        "__fields_types__": fields_types,
        "__size__": size,
        "__is_struct__": not __is_union__,
        "__is_union__": __is_union__,
        "__is_enum__": False,
        "__byte_order__": __byte_order__,
        "__alignment__": max_alignment,
        "__name__": __name__,
        "__cls__": __cls__,
    }
    return result


import cv2
import os
import numpy as np
import random

# normalize each image by dividing the pixel values with the
# standard deviation of the pixel values of the images in a dataset

# def resize_img(path):
# 	# threshold = 220
# 	try:
# 		img = cv2.imread('CEDAR/full_forg/' + path, 0)
# 		dst = cv2.resize(img, (220, 155), cv2.INTER_LINEAR)
# 		cv2.imwrite('CEDAR/full_forg_gray_115x220/{}'.format(path), dst)
# 	except:
# 		print(path)

# path = 'CEDAR/full_forg'
# for p in os.listdir(path):
# 	resize_img(p)

# img = cv2.imread('dataset/original_2_9.png', 0)
# # img_ = cv2.resize(img, (img.shape[0]*2, img.shape[1]*2), cv2.INTER_NEAREST)
# print(img.shape)
# cv2.imshow('origin', img)
# cv2.imshow('resize', 255 - img)

# cv2.waitKey()

with open('CEDAR/gray_train.txt', 'w') as f:
	for i in range(1, 51):
		for j in range(1, 25):
			for k in range(j+1, 25):
				f.write('full_org_gray_115x220/original_{0}_{1}.png full_org_gray_115x220/original_{0}_{2}.png 1\n'.format(i, j, k))
		org_forg = [(j,k) for j in range(1, 25) for k in range(1, 25)]
		for (j, k) in random.choices(org_forg, k=276):
			f.write('full_org_gray_115x220/original_{0}_{1}.png full_forg_gray_115x220/forgeries_{0}_{2}.png 0\n'.format(i, j, k))

with open('CEDAR/gray_test.txt', 'w') as f:
	for i in range(51, 56):
		for j in range(1, 25):
			for k in range(j+1, 25):
				f.write('full_org_gray_115x220/original_{0}_{1}.png full_org_gray_115x220/original_{0}_{2}.png 1\n'.format(i, j, k))
		org_forg = [(j,k) for j in range(1, 25) for k in range(1, 25)]
		for (j, k) in random.choices(org_forg, k=276):
			f.write('full_org_gray_115x220/original_{0}_{1}.png full_forg_gray_115x220/forgeries_{0}_{2}.png 0\n'.format(i, j, k))


# A = input()

# for i in range(len(A)):
#     for j in range(len(A)):
#         for k in range(len(A)):

a = input()
lst = []
for i in range(1, len(a)):       # 모든 경우의 수를 구한다
    for j in range(i+1, len(a)):
        first = a[:i][::-1]   # m
        second = a[i:j][::-1] # o
        third = a[j:][::-1] # letib
        lst.append(first + second + third) # 글자수를 유지하며 반복한 모든 수를 더한다
lst.sort()
print(lst[0]) #  정렬 시 정답이 알파벳 순서상 첫번째로 나오기 때문에 [0]을 통해 출력


import unittest

import w3af.core.data.kb.knowledge_base as kb
from w3af.core.data.url.HTTPResponse import HTTPResponse
from w3af.core.data.request.fuzzable_request import FuzzableRequest
from w3af.core.data.parsers.doc.url import URL
from w3af.core.data.dc.headers import Headers
from w3af.core.controllers.misc.temp_dir import create_temp_dir
from w3af.plugins.grep.wsdl_greper import wsdl_greper

class test_wsdl_greper(unittest.TestCase):

    def setUp(self):
        create_temp_dir()
        kb.kb.cleanup()
        self.plugin = wsdl_greper()
        self.url = URL('http://www.w3af.com/')
        self.request = FuzzableRequest(self.url)

    def tearDown(self):
        self.plugin.end()

    def test_wsdl_greper_empty(self):
        body = ''
        headers = Headers([('content-type', 'text/html')])
        response = HTTPResponse(200, body, headers, self.url, self.url, _id=1)
        self.plugin.grep(self.request, response)
        self.assertEqual(len(kb.kb.get('wsdl_greper', 'wsdl')), 0)

    def test_wsdl_greper_long(self):
        body = 'ABC ' * 10000
        headers = Headers([('content-type', 'text/html')])
        response = HTTPResponse(200, body, headers, self.url, self.url, _id=1)
        self.plugin.grep(self.request, response)
        self.assertEqual(len(kb.kb.get('wsdl_greper', 'wsdl')), 0)

    def test_wsdl_greper_positive(self):
        body = 'ABC ' * 100
        body += '/s:sequence'
        body += '
 ' * 50
        headers = Headers([('content-type', 'text/html')])
        response = HTTPResponse(200, body, headers, self.url, self.url, _id=1)
        self.plugin.grep(self.request, response)
        self.assertEqual(len(kb.kb.get('wsdl_greper', 'wsdl')), 1)

    def test_wsdl_greper_positive_disco(self):
        body = 'ABC ' * 100
        body += 'disco:discovery '
        body += '
 ' * 50
        headers = Headers([('content-type', 'text/html')])
        response = HTTPResponse(200, body, headers, self.url, self.url, _id=1)
        self.plugin.grep(self.request, response)
        self.assertEqual(len(kb.kb.get('wsdl_greper', 'disco')), 1)
        self.assertEqual(len(kb.kb.get('wsdl_greper', 'wsdl')), 0)


N = int(input())
cnt = 0
Check = N
while True:
  Temp = N//10 + N%10
  newN = (N%10)*10 + Temp%10
  cnt += 1
  N=newN
  if newN == Check:
    break
print(cnt)

# 각 자리 숫자를 추출하는 법 => 10을 나눈 몫


import matplotlib.pyplot as plt
import mlflow
import pandas as pd
from ml_gapfill.gsmm_analysis import GSMMAnalysis

import plotly
import plotly.express as px
import plotly.graph_objects as go

import numpy as np
from sklearn.metrics import ConfusionMatrixDisplay

import pickle
import plotly.figure_factory as ff

colours = [
    "#003f5c",
    "#58508d",
    "#bc5090",
    "#ff6361",
    "#ffa600",
    "#ffa800",
]

def plot_train_pedict_scatter(y_train, y_pred, output_path):
    fig, ax = plt.subplots()
    ax.scatter(y_train, y_pred)
    plt.xlabel("label")
    plt.ylabel("prediction")
    mlflow.log_figure(fig, output_path)
    plt.close()

def plot_prediction_label_scatter(model, data_df, label_column="Y_OD"):
    feature_names = data_df.columns.values[data_df.columns.str.startswith("X_")]

    X_data = data_df.loc[:, data_df.columns.str.startswith("X_")].values
    y_label = data_df.loc[:, label_column].values
    y_pred = model.predict(X_data)

    layout_min = np.min([np.min(y_label), np.min(y_pred)]) - 0.1
    layout_max = np.max([np.max(y_label), np.max(y_pred)]) + 0.1

    fig = go.Figure(
        data=go.Scatter(
            x=y_label,
            y=y_pred,
            marker={"color": colours[2]},
            mode="markers",
            opacity=0.75,
        ),
        layout_yaxis_range=[layout_min, layout_max],
        layout_xaxis_range=[layout_min, layout_max],
    )

    fig.update_layout(xaxis_title="Actual", yaxis_title="Predicted")
    fig.update_layout(template="simple_white", width=1000, height=1000)

    return fig

def plot_confusion_matrix(y_test, X_test, model, output_path):
    # Plot test confusion matrix
    disp = ConfusionMatrixDisplay.from_estimator(
        model,
        X_test,
        y_test,
        cmap=plt.cm.Blues,
    )

    mlflow.log_figure(disp.figure_, output_path)

def plotly_fig2array(fig):
    # convert Plotly fig to  an array
    fig_bytes = fig.to_image(format="pdf")
    buf = io.BytesIO(fig_bytes)
    img = Image.open(buf)
    return np.asarray(img)

def plot_models_dendrogram(
    df,
    output_path,
    models_dir="/Users/bezk/Documents/CAM/research_code/community_study_mel_krp/analysis_daniel_machado/models/",
    dendro_figure_path="/Users/bezk/Documents/CAM/research_code/community_study_mel_krp/ml_project/scripts/preprocessing_scripts/dendrogram_figure.pkl",
    confusion=True,
):

    # Load figure template
    with open(dendro_figure_path, "rb") as f:
        fig = pickle.load(f)

    dendro_leaves = fig["layout"]["xaxis"]["ticktext"]

    # Make column to plot
    # nan=0, True negative=1, False positive=2, False negative=3, True positive=4
    def confusion_label_gen(row):
        if pd.isna(row["Y_pred"]):
            return 0

        # True negative
        elif row["Y_pred"] == 0 and row["Y_growth_class"] == 0:
            return 1

        # False positive
        elif row["Y_pred"] == 1 and row["Y_growth_class"] == 0:
            return 2

        # False negative
        elif row["Y_pred"] == 0 and row["Y_growth_class"] == 1:
            return 3

        # True positive
        elif row["Y_pred"] == 1 and row["Y_growth_class"] == 1:
            return 4

    def accuracy_label_gen(row):
        if pd.isna(row["Y_pred"]):
            print("isnan")
            return 0

        elif row["Y_pred"] == 0 and row["Y_growth_class"] == 0:
            return 2

        elif row["Y_pred"] == 1 and row["Y_growth_class"] == 1:
            return 2

        else:
            return 1

    if confusion:
        df["confusion_label"] = df.apply(confusion_label_gen, axis=1)

        df = df.pivot(
            index="Y_media_label", columns="Y_species_label", values="confusion_label"
        )
        colorscale = [
            [0.0, "rgb(0,0,0)"],
            [0.25, "rgb(9, 118, 148)"],
            [0.5, "rgb(209, 122, 23)"],
            [0.75, "rgb(148, 9, 39)"],
            [1.0, "rgb(9, 148, 44)"],
        ]

    else:
        df["accuracy_label"] = df.apply(accuracy_label_gen, axis=1)

        df = df.pivot(
            index="Y_media_label", columns="Y_species_label", values="accuracy_label"
        )
        colorscale = [
            [0.0, "rgb(0,0,0)"],
            [0.5, "rgb(148, 9, 39)"],
            [1.0, "rgb(9, 148, 44)"],
        ]

    # Make new column with media only digits
    df["media_only_digits"] = df.index.str.extract(r"(\d+)").astype(int).values

    # Sort by digit column
    df.sort_values(by=["media_only_digits"], inplace=True)

    # Drop column
    df.drop(columns=["media_only_digits"], inplace=True)

    # dendro_leaves = [d for d in dendro_leaves if d in df.columns]

    all_model_names = dendro_leaves
    for name in all_model_names:
        if name not in df.columns:
            df[name] = np.nan

    n_rows, n_cols = df.shape

    heatmap_vals = df[dendro_leaves].values

    heatmap_vals[pd.isna(heatmap_vals)] = 0
    # Set anything above zero to one
    # heatmap_vals[heatmap_vals > 0] = 1

    heatmap = [
        go.Heatmap(
            x=list(range(len(dendro_leaves))),
            y=list(range(n_rows)),
            z=heatmap_vals,
            colorscale=colorscale,
        )
    ]

    heatmap[0]["x"] = dendro_leaves
    # Add Heatmap Data to Figure
    for data in heatmap:
        fig.add_trace(data)

    # Edit xaxis
    fig.update_layout(
        xaxis={
            "domain": [0.15, 1],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "ticks": "",
        }
    )
    # Edit xaxis2
    fig.update_layout(
        xaxis2={
            "domain": [0, 0.15],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "showticklabels": False,
            "ticks": "",
        }
    )

    # Edit yaxis
    fig.update_layout(
        yaxis={
            "domain": [0, 0.85],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "showticklabels": True,
            "tickvals": [x - 0.25 for x in list(range(n_rows))],
            "ticktext": list(df.index),
        }
    )
    # Edit yaxis2
    fig.update_layout(
        yaxis2={
            "domain": [0.825, 0.975],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "showticklabels": False,
            "ticks": "",
        }
    )

    fig.update_layout(width=1000, height=1000, showlegend=False)
    mlflow.log_figure(fig, output_path)

    # fig.show()
    # data_dist = pdist(data_array)
    # heat_data = squareform(data_dist)
    # heat_data = heat_data[dendro_leaves,:]
    # heat_data = heat_data[:,dendro_leaves]

    return fig

def plot_models_dendrogram_v2(df, output_path, confusion=True):
    # Extract model features by subsetting for species and keeping constant columns
    unique_species = df.Y_species_label.unique()
    unique_medias = df.Y_media_label.unique()
    model_features = []

    for species in unique_species:
        sub_df = df.loc[df["Y_species_label"] == species]
        features_df = sub_df.loc[:, (sub_df == sub_df.iloc[0]).any()]

        select_columns = list(
            features_df.columns.values[features_df.columns.str.startswith("X_")]
        )

        features_df = sub_df[select_columns]
        features = features_df.values[0]
        model_features.append(features)

    model_features = np.array(model_features)

    fig = ff.create_dendrogram(
        model_features, orientation="bottom", labels=unique_species
    )

    fig.update_layout(width=800, height=800, showlegend=False)

    for i in range(len(fig["data"])):
        fig["data"][i]["yaxis"] = "y2"

    dendro_leaves = fig["layout"]["xaxis"]["ticktext"]

    # Make column to plot
    # nan=0, True negative=1, False positive=2, False negative=3, True positive=4
    def confusion_label_gen(row):
        if pd.isna(row["Y_pred"]):
            return 0

        # True negative
        elif row["Y_pred"] == 0 and row["Y_growth_class"] == 0:
            return 1

        # False positive
        elif row["Y_pred"] == 1 and row["Y_growth_class"] == 0:
            return 2

        # False negative
        elif row["Y_pred"] == 0 and row["Y_growth_class"] == 1:
            return 3

        # True positive
        elif row["Y_pred"] == 1 and row["Y_growth_class"] == 1:
            return 4

    def accuracy_label_gen(row):
        if pd.isna(row["Y_pred"]):
            print("isnan")
            return 0

        elif row["Y_pred"] == 0 and row["Y_growth_class"] == 0:
            return 2

        elif row["Y_pred"] == 1 and row["Y_growth_class"] == 1:
            return 2

        else:
            return 1

    df["accuracy_label"] = df.apply(accuracy_label_gen, axis=1)

    if confusion:
        df["confusion_label"] = df.apply(confusion_label_gen, axis=1)

        df = df.pivot(
            index="Y_media_label", columns="Y_species_label", values="confusion_label"
        )
        colorscale = [
            # [0.0, "rgb(0,0,0)"],
            [0.0, "rgb(9, 148, 44)"],
            [0.333, "rgb(255, 255, 102)"],
            [0.666, "rgb(255, 128, 39)"],
            [1.0, "rgb(9, 148, 44)"],
        ]

    else:
        df = df.pivot(
            index="Y_media_label", columns="Y_species_label", values="accuracy_label"
        )
        colorscale = [
            [0.0, "rgb(148, 9, 39)"],
            [1.0, "rgb(9, 148, 44)"],
        ]

    # Make new column with media only digits
    df["media_only_digits"] = df.index.str.extract(r"(\d+)").astype(int).values

    # Sort by digit column
    df.sort_values(by=["media_only_digits"], inplace=True)

    # Drop column
    df.drop(columns=["media_only_digits"], inplace=True)
    all_model_names = dendro_leaves
    for name in all_model_names:
        if name not in df.columns:
            df[name] = np.nan

    n_rows, n_cols = df.shape

    heatmap_vals = df[dendro_leaves].values

    heatmap_vals[pd.isna(heatmap_vals)] = 0
    # Set anything above zero to one
    # heatmap_vals[heatmap_vals > 0] = 1

    heatmap = [
        go.Heatmap(
            x=list(range(len(dendro_leaves))),
            y=list(range(n_rows)),
            z=heatmap_vals,
            colorscale=colorscale,
        )
    ]

    heatmap[0]["x"] = fig["layout"]["xaxis"]["tickvals"]
    # fig.update_layout(width=1500, height=1000, showlegend=False)

    # Add Heatmap Data to Figure
    for data in heatmap:
        fig.add_trace(data)

    # Edit xaxis
    fig.update_layout(
        xaxis={
            "domain": [0.15, 1],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "tickangle": -45,
            "ticks": "",
        }
    )
    # Edit xaxis2
    fig.update_layout(
        xaxis2={
            "domain": [0, 0.15],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "showticklabels": False,
            "tickangle": -45,
            "ticks": "",
        }
    )

    fig.update_xaxes(tickfont=dict(size=12))

    # Edit yaxis
    fig.update_layout(
        yaxis={
            "domain": [0, 0.85],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "showticklabels": True,
            "tickvals": [x - 0.25 for x in list(range(n_rows))],
            "ticktext": list(df.index),
        }
    )
    # Edit yaxis2
    fig.update_layout(
        yaxis2={
            "domain": [0.85, 0.975],
            "mirror": False,
            "showgrid": False,
            "showline": False,
            "zeroline": False,
            "showticklabels": False,
            "ticks": "",
        }
    )

    return fig

if __name__ == "__main__":
    wd = "/Users/bezk/Documents/CAM/research_code/community_study_mel_krp"
    models_dir = f"{wd}/analysis_daniel_machado/models/"
    df = pd.read_csv(
        "/Users/bezk/Documents/CAM/research_code/community_study_mel_krp/ml_project/mlruns/1/7fb276c3b78246c9871985d5a916483f/artifacts/test_df.csv"
    )

    plot_models_dendrogram(df, output_path="./test.html")


import unittest
from app import db, create_app, models, utils
from app.api import db_views
import json
from datetime import datetime, timedelta
from tests.fake_db import gen_fake

class DBViewTestCase(unittest.TestCase):
    def setUp(self):
        self.app = create_app('testing')
        self.app_context = self.app.app_context()
        self.app_context.push()
        db.create_all()
        self.client = self.app.test_client()
        gen_fake()

    def tearDown(self):
        db.session.remove()
        db.drop_all()
        self.app_context.pop()

    def test_db_view_api_ping(self):
        # check the avaliablity of apis.
        responses = [self.client.get('/'),
                     self.client.get('/api/v1/project/all'),
                     self.client.get('/api/v1/project/1/spots'),
                     self.client.get('/api/v1/device/1/records'),
                     self.client.get('/api/v1/project_pic/1')]
        self.assertTrue(all(filter(lambda r: r.status_code == 200, responses)))

    # def test_spot_detailed_view_outdoor_data_integrity(self):
    #     dday = timedelta(days=1)
    #     dhour = timedelta(hours=1)
    #     days_in_test_month = [datetime(2019, 5, 1, 0) + i * dday + h * dhour
    #                           for h in range(0, 24)
    #                           for i in range(0, 31)]

    #     cheap_gen_fake()
    #     # generate test outdoor data
    #     for day in days_in_test_month:
    #         odr = models.OutdoorRecord(outdoor_record_time=day)
    #         db.session.add(odr)
    #         db.session.commit()

    #     models.rand_date = utils.rand_date_in(datetime(2019, 5, 1),
    #                                           datetime(2019, 6, 1))
    #     models.SpotRecord.gen_fake(500)

    #     response = self.client.get('/view/spot/1')
    #     response_data = json.loads(response.get_data(as_text=True))

    #     # if outdoor data exists in the corresponding records.
    #     has_outdoor = any(filter(lambda data: data['outdoor_record'] != {}, response_data))
    #     self.assertTrue(has_outdoor)

    #     # if outdoor data hour is correct.
    #     def is_in_same_hour(d1, d2) -> bool:
    #         _, day1, mon1, year1, time1, _ = d1.split()
    #         _, day2, mon2, year2, time2, _ = d2.split()
    #         hour1 = time1.split(":")[0]
    #         hour2 = time2.split(":")[0]
    #         d1_str = day1 + mon1 + year1 + hour1
    #         d2_str = day2 + mon2 + year2 + hour2
    #         return d1_str == d2_str

    #     same_hour = False
    #     for data in response_data:
    #         od_time: str = data["outdoor_record"]["outdoor_record_time"]
    #         rec_time: str = data["spot_record_time"]
    #         if is_in_same_hour(od_time, rec_time):
    #             same_hour = True
    #     self.assertTrue(same_hour)