"""
This file contains the classes and functions to simulate the propagation of light.
"""
from __future__ import annotations

import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import copy
import dataclasses
import itertools
from typing import List, Tuple, Union

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from matplotlib.ticker import EngFormatter
from scipy import interpolate

from metabox import expansion, utils

@dataclasses.dataclass
class FieldProperties:
    """Defines the properties of a 2d field.

    Args:
        n_pixels: the number of pixels per dim after 2D expansion.
        wavelength: the wavelength of the light in meters.
        theta: the angle of the light in degrees.
        phi: the phase of the light in degrees.
        period: the period of the pixels in meters.
        upsampling: the upsampling factor.
        use_padding: whether to use padding or not.
        use_antialiasing: whether to use antialiasing or not.
    """

    n_pixels: int
    wavelength: List[float]
    theta: List[float]
    phi: List[float]
    period: float
    upsampling: int
    use_padding: bool
    use_antialiasing: bool

    def copy(self):
        """Returns a copy of the field properties."""
        return copy.deepcopy(self)

@dataclasses.dataclass
class Field2D(FieldProperties):
    """Class to store the field data and its metadata.

    Args:
        tensor: the tensor of the field.
        wavelength: the wavelength of the light in meters.
        theta: the angle of the light in degrees.
        phi: the phase of the light in degrees.
        period: the period of the pixels in meters.
        upsampling: the upsampling factor.
        use_padding: whether to use padding or not.
        use_antialiasing: whether to use antialiasing or not.
    """

    tensor: tf.Tensor

    def modulated_by(self, other: Field2D) -> Field2D:
        """Modulate this field by another field.

        Args:
            other (Field2D): the field to modulate by.

        Returns:
            Field2D: the modulated field.
        """
        # Check that the fields have the same properties.
        if not np.all(self.wavelength == other.wavelength):
            raise ValueError(
                "Wavelengths are not the same, got  {0} and {1}.".format(
                    self.wavelength, other.wavelength
                )
            )
        if not self.period == other.period:
            if self.period > other.period:
                fine_field = other.copy()
                coarse_field = self.copy()
            else:
                fine_field = self.copy()
                coarse_field = other.copy()
            enlarge_factor = coarse_field.period / float(fine_field.period)
            old_shape = coarse_field.tensor.shape[-2:]
            new_shape = (
                int(old_shape[0] * enlarge_factor),
                int(old_shape[1] * enlarge_factor),
            )
            real_part = tf.math.real(coarse_field.tensor)
            imag_part = tf.math.imag(coarse_field.tensor)

            parts = []
            for part in [real_part, imag_part]:
                part = tf.image.resize(
                    part[..., tf.newaxis],
                    new_shape,
                    method=tf.image.ResizeMethod.NEAREST_NEIGHBOR,
                )
                part = tf.image.resize_with_crop_or_pad(
                    part,
                    fine_field.tensor.shape[-2],
                    fine_field.tensor.shape[-1],
                )
                part = part[..., 0]
                parts.append(part)
            mod_tensor = tf.complex(parts[0], parts[1])

            self = fine_field
            other = copy.deepcopy(fine_field)
            other.tensor = mod_tensor
        if not np.all(self.phi == other.phi):
            raise ValueError("angles_y are not the same.")
        if not np.all(self.theta == other.theta):
            raise ValueError("Angles are not the same.")
        if not self.tensor.shape == other.tensor.shape:
            if self.tensor.shape[-2] > other.tensor.shape[-2]:
                small_field = other
                large_field = self
            else:
                small_field = self
                large_field = other

            real_part = tf.math.real(small_field.tensor)
            imag_part = tf.math.imag(small_field.tensor)
            parts = []
            for part in [real_part, imag_part]:
                part = tf.image.resize_with_crop_or_pad(
                    part[..., tf.newaxis],
                    large_field.tensor.shape[-2],
                    large_field.tensor.shape[-1],
                )
                part = part[..., 0]
                parts.append(part)
            mod_tensor = tf.complex(parts[0], parts[1])

            self = large_field
            other = copy.deepcopy(large_field)
            other.tensor = mod_tensor

        new_field = copy.deepcopy(self)
        new_field.tensor *= other.tensor
        return new_field

    def __post_init__(self):
        # Check that the field tensor is of the correct shape.
        if len(self.tensor.shape) != 3:
            raise ValueError(
                "Field tensor must have shape [batchsize, n_pix, n_pix]"
            )

        # check batch size is correct
        n_batch = self.tensor.shape[0]
        expected_n_batch = (
            len(self.wavelength) * len(self.theta) * len(self.phi)
        )
        if n_batch != expected_n_batch:
            raise ValueError(
                """Batch size of field tensor does not match the number of
                wavelengths, angles, and angles_y multiplied together.
                Expected batch size: {0}
                Got batch size: {1}""".format(
                    expected_n_batch, n_batch
                )
            )

    def get_intensity(self):
        """Returns the intensity tensor of the tensor."""
        return get_intensity_2d(self)

    def get_phase(self):
        """Returns the phase tensor of the tensor."""
        return get_phase_2d(self)

    def wavelength_average(self):
        """Returns the wavelength averaged field."""
        return wavelength_average_2d(self)

    def show_phase(self):
        """Shows the phase of the field."""

        phase = self.get_phase()
        wl, ag = phase.shape[:2]
        diameter = self.period * phase.shape[-1]
        radius = diameter / 2.0
        for i, j in itertools.product(range(wl), range(ag)):
            wave = self.wavelength[i] * 1e6
            angle = list(itertools.product(self.theta, self.phi))[j]
            f = plt.figure(figsize=(5, 5), dpi=100)
            ax = plt.axes([0, 0.05, 0.9, 0.9])
            im = ax.imshow(
                phase[i, j], extent=[-radius, radius, -radius, radius]
            )
            formatter0 = EngFormatter(unit="m")
            ax.xaxis.set_major_formatter(formatter0)
            ax.yaxis.set_major_formatter(formatter0)
            plt.locator_params(axis="y", nbins=3)
            plt.locator_params(axis="x", nbins=3)
            ax.set_xlabel("X")
            ax.set_ylabel("Y")
            ax.grid(False)
            title = "Phase Distribution\n λ={0}µm, AOI={1}°,{2}°".format(
                round(wave, 2), round(angle[0], 2), round(angle[1], 2)
            )
            ax.set_title(title)
            cax = plt.axes([0.95, 0.05, 0.05, 0.9])
            plt.colorbar(mappable=im, cax=cax)
            plt.show()

    def show_intensity(self, crop_factor=1.0):
        """Shows the intensity of the field.

        Args:
            crop_factor (float): The crop factor. Must be less than or equal to 1.0.
        """

        if crop_factor > 1.0:
            raise ValueError("Zoom must be greater than or equal to 1.0.")

        intensity = self.get_intensity()
        # put angle in the last dimension
        intensity = tf.transpose(intensity, [0, 2, 3, 1])
        if crop_factor != 1.0:
            intensity = tf.image.central_crop(intensity, crop_factor)
        # return to the original dim order
        intensity = tf.transpose(intensity, [0, 3, 1, 2])
        wl, ag = intensity.shape[:2]
        diameter = self.period * intensity.shape[-1] * crop_factor
        radius = diameter / 2.0
        for i, j in itertools.product(range(wl), range(ag)):
            wave = self.wavelength[i] * 1e6
            angle = list(itertools.product(self.theta, self.phi))[j]
            f = plt.figure(figsize=(5, 5), dpi=100)
            ax = plt.axes([0, 0.05, 0.9, 0.9])
            im = ax.imshow(
                intensity[i, j], extent=[-radius, radius, -radius, radius]
            )
            formatter0 = EngFormatter(unit="m")
            ax.xaxis.set_major_formatter(formatter0)
            ax.yaxis.set_major_formatter(formatter0)
            plt.locator_params(axis="y", nbins=3)
            plt.locator_params(axis="x", nbins=3)
            ax.set_xlabel("X")
            ax.set_ylabel("Y")
            ax.grid(False)
            title = "Intensity Distribution\n λ={0}µm, AOI={1}°,{2}°".format(
                round(wave, 2), round(angle[0], 2), round(angle[1], 2)
            )
            ax.set_title(title)
            cax = plt.axes([0.95, 0.05, 0.05, 0.9])
            plt.colorbar(mappable=im, cax=cax)
            plt.show()

    def show_color_intensity(self, crop_factor=1.0):
        """Shows the intensity of the field.

        Args:
            crop_factor (float): The crop factor. Must be less than or equal to 1.0.
        """

        if crop_factor > 1.0:
            raise ValueError("Zoom must be greater than or equal to 1.0.")

        rgb_intensity = self.to_rgb_intensity()
        if crop_factor != 1.0:
            rgb_intensity = tf.image.central_crop(rgb_intensity, crop_factor)
        ag = rgb_intensity.shape[0]
        diameter = self.period * rgb_intensity.shape[-2] * crop_factor
        radius = diameter / 2.0
        for j in range(ag):
            if len(self.wavelength) == 1:
                wave = round(self.wavelength[0] * 1e6, 2)
            else:
                wave_0 = round(self.wavelength[0] * 1e6, 2)
                wave_1 = round(self.wavelength[-1] * 1e6, 2)
                wave = f"{wave_0}-{wave_1}"
            angle = list(itertools.product(self.theta, self.phi))[j]
            f = plt.figure(figsize=(5, 5), dpi=100)
            ax = plt.axes([0, 0.05, 0.9, 0.9])
            im = ax.imshow(
                rgb_intensity[j], extent=[-radius, radius, -radius, radius]
            )
            formatter0 = EngFormatter(unit="m")
            ax.xaxis.set_major_formatter(formatter0)
            ax.yaxis.set_major_formatter(formatter0)
            plt.locator_params(axis="y", nbins=3)
            plt.locator_params(axis="x", nbins=3)
            ax.set_xlabel("X")
            ax.set_ylabel("Y")
            ax.grid(False)
            title = (
                "Color Intensity Distribution\n λ={0}µm, AOI={1}°,{2}°".format(
                    wave, round(angle[0], 2), round(angle[1], 2)
                )
            )
            ax.set_title(title)
            plt.show()

    def to_rgb_intensity(self):
        """Return RGB image of the intensity."""

        intensity_distributions = self.get_intensity().numpy()
        weighted_intensity_distributions = []

        for wavelength, intensity_dist in zip(
            self.wavelength, intensity_distributions
        ):
            r, g, b = utils.wavelength_to_rgb(wavelength)
            r_dist = intensity_dist * r
            g_dist = intensity_dist * g
            b_dist = intensity_dist * b
            rgb_dist = np.stack([r_dist, g_dist, b_dist], axis=-1)
            weighted_intensity_distributions.append(rgb_dist)
        # normalize the image
        weighted_intensity_distributions = np.array(
            weighted_intensity_distributions
        )
        weighted_intensity_distributions /= np.max(
            weighted_intensity_distributions
        )
        return np.sum(weighted_intensity_distributions, axis=0)

@dataclasses.dataclass
class Field1D(FieldProperties):
    """Class to store the field data and its metadata.

    Args:
        wavelength: the wavelength of the light in meters.
        theta: the angle of the light in degrees.
        phi: the phase of the light in degrees.
        period: the period of the pixels in meters.
        upsampling: the upsampling factor.
        use_padding: whether to use padding or not.
        use_antialiasing: whether to use antialiasing or not.
    """

    tensor: tf.Tensor

    def __post_init__(self):
        # Check that the field tensor is of the correct shape.
        if len(self.tensor.shape) != 2:
            raise ValueError(
                "Field tensor must have shape [batchsize, n_pix_radius]"
            )

        # check batch size is correct
        n_batch = self.tensor.shape[0]
        expected_n_batch = (
            len(self.wavelength) * len(self.theta) * len(self.phi)
        )
        if n_batch != expected_n_batch:
            raise ValueError(
                """Batch size of field tensor does not match the number of
                wavelengths, angles, and angles_y multiplied together."""
            )

    def expand_to_2d(self, basis_dir="basis_data") -> Field2D:
        """Function to expand a 1d field to a 2d field.

        Args:
            basis_dir: the directory where the basis is saved.
                The default directory is "basis_data".

        Returns:
            propagation2d.Field2D: a 2d field
        """
        new_tensor = expansion.expand_to_2d(self.tensor, basis_dir)

        # create the 2d field
        field2d = Field2D(
            tensor=new_tensor,
            n_pixels=self.n_pixels,
            wavelength=self.wavelength,
            theta=self.theta,
            phi=self.phi,
            period=self.period,
            upsampling=self.upsampling,
            use_padding=self.use_padding,
            use_antialiasing=self.use_antialiasing,
        )

        return field2d

    def get_intensity(self):
        return get_intensity_1d(self)

def get_intensity_1d(field_1d: Field1D):
    """returns the intensity tensor of the field.

    Returns:
        tf.Tensor: intensity tensor of shape [wavelengths, angles, pixelsX]
    """

    intensity = tf.math.abs(field_1d.tensor) ** 2

    # new shape
    new_shape = tf.convert_to_tensor(
        [
            len(field_1d.wavelength),
            len(field_1d.theta) * len(field_1d.phi),
            field_1d.tensor.shape[1],
        ]
    )
    return tf.reshape(intensity, new_shape)

def get_phase_2d(field_2d: Field2D):
    """returns the phase tensor of the field.

    Returns:
        tf.Tensor: intensity tensor of shape [wavelengths, angles, pixelsX, pixelsY]
    """

    phase = tf.math.angle(field_2d.tensor)

    # new shape
    new_shape = tf.convert_to_tensor(
        [
            len(field_2d.wavelength),
            len(field_2d.theta) * len(field_2d.phi),
            field_2d.tensor.shape[1],
            field_2d.tensor.shape[2],
        ]
    )
    return tf.reshape(phase, new_shape)

def get_intensity_2d(field_2d: Field2D):
    """returns the intensity tensor of the field.

    Returns:
        tf.Tensor: intensity tensor of shape [wavelengths, angles, pixelsX, pixelsY]
    """

    intensity = tf.math.abs(field_2d.tensor) ** 2

    # new shape
    new_shape = tf.convert_to_tensor(
        [
            len(field_2d.wavelength),
            len(field_2d.theta) * len(field_2d.phi),
            field_2d.tensor.shape[1],
            field_2d.tensor.shape[2],
        ]
    )
    return tf.reshape(intensity, new_shape)

def wavelength_average_2d(field: Field2D) -> Field2D:
    """Function to average the field over the wavelengths.

    Args:
        field (Field2D): the field to average over the wavelengths.

    Returns:
        Field2D: the averaged field.
    """

    # seperate wavelength and angle dims
    field_tensor = tf.reshape(
        field.tensor,
        [
            len(field.wavelength),
            len(field.theta) * len(field.phi),
            field.n_pixels,
            field.n_pixels,
        ],
    )

    # average over the wavelengths
    new_tensor = tf.reduce_mean(field_tensor, axis=0, keepdims=False)

    # create the new field
    new_field = Field2D(
        tensor=new_tensor,
        n_pixels=field.n_pixels,
        wavelength=[np.mean(field.wavelength)],
        theta=field.theta,
        phi=field.phi,
        period=field.period,
        upsampling=field.upsampling,
        use_padding=field.use_padding,
        use_antialiasing=field.use_antialiasing,
    )

    return new_field

def get_transfer_function(
    field_like: Field2D,
    ref_idx: float,
    prop_dist: float,
    lateral_shift: Union[None, Tuple[float, float]] = None,
) -> tf.Tensor:
    """Get the Propagator object given the propagation information.

    Args:
        ref_idx (float): refractive index of the medium
        prop_dist (float): propagation distance in meters
        upsampling (int): upsampling factor
        use_padding (bool): whether or not to use padding
        use_anti_aliasing (bool): whether or not to use anti-aliasing
        lateral_shift: the lateral shift of the sampling window on the detector
            in meters. If None, the shift is set so that the Chief Ray is at the
            center of the detector. If a tuple of two floats, the shift is set
            according to the first element (x shift) and the second element (y
            shift) of the input tuple.

    Returns:
        tf.Tensor: the complex field on the final plane
    """

    return get_propagator_batched(
        ref_idx=ref_idx,
        prop_dist=prop_dist,
        n_pix=field_like.tensor.shape[1],
        period=field_like.period,
        wavelength_sampling=field_like.wavelength,
        theta_sampling=field_like.theta,
        phi_sampling=field_like.phi,
        upsampling=field_like.upsampling,
        use_padding=field_like.use_padding,
        use_antialiasing=field_like.use_antialiasing,
        lateral_shift=lateral_shift,
    )

def propagate(
    field: Field2D,
    transfer_function: tf.Tensor,
) -> Field2D:
    """Propagate a field through a given transfer function.

    Args:
        field (Field): the field to propagate
        transfer_function (tf.Tensor): the transfer function

    Returns:
        End field (Field): the complex field on the final plane
    """
    field = copy.deepcopy(field)
    new_tensor = propagate_with_propagator_batched(
        field.tensor,
        transfer_function,
        upsampling=field.upsampling,
        use_padding=field.use_padding,
    )
    field.tensor = new_tensor
    return field

def get_propagator_batched(
    ref_idx: float,
    prop_dist: float,
    n_pix: int,
    period: float,
    wavelength_sampling: List[float],
    theta_sampling: List[float],
    phi_sampling: List[float],
    upsampling=1,
    use_padding=True,
    use_antialiasing=True,
    lateral_shift: Union[None, Tuple[float, float]] = None,
) -> tf.Tensor:
    """
    Returns the transfer function for a given propagation distance.

    Args:
        ref_idx: refractive index of the medium
        prop_dist: propagation distance in meters
        n_pix: number of pixels in each dimension
        period: pixel size in meters
        wavelength: wavelength in meters
        theta_sampling: list of angles to sample in degrees
        phi_sampling: list of angles to sample in degrees
        upsampling: upsampling factor
        use_padding: whether to pad the transfer function to prevent aliasing
        use_antialiasing: whether to limit the transfer function bandwidth to
            prevent aliasing
        lateral_shift: the lateral shift of the sampling window on the detector
            in meters. If None, the shift is set so that the Chief Ray is at the
            center of the detector. If a tuple of two floats, the shift is set
            according to the first element (x shift) and the second element (y
            shift) of the input tuple.

    Returns:
        propagator: transfer function
    """
    batch_size = (
        len(wavelength_sampling) * len(theta_sampling) * len(phi_sampling)
    )

    lam0 = tf.convert_to_tensor(
        np.repeat(
            wavelength_sampling,
            np.size(theta_sampling) * np.size(phi_sampling),
        ),
        dtype=tf.float32,
    )
    lam0 = lam0[:, tf.newaxis, tf.newaxis]
    lam0 = tf.tile(lam0, multiples=(1, n_pix, n_pix))

    # Propagator definition.
    k = ref_idx * 2 * np.pi / lam0[:, 0, 0]
    k = k[:, np.newaxis, np.newaxis]
    samp = int(upsampling * n_pix)
    k = tf.cast(k, dtype=tf.complex64)

    if use_padding:
        k_xlist_pos = (
            2 * np.pi * np.linspace(0, 1 / (2 * period / upsampling), samp)
        )
        front = k_xlist_pos[-(samp - 1) :]
        front = -front[::-1]
        k_xlist = np.hstack((front, k_xlist_pos))
        k_x = np.kron(k_xlist, np.ones((2 * samp - 1, 1)))
    else:
        k_xlist = (
            2
            * np.pi
            * np.linspace(
                -1 / (2 * period / upsampling),
                1 / (2 * period / upsampling),
                samp,
            )
        )
        k_x = np.kron(k_xlist, np.ones((samp, 1)))

    k_x = k_x[np.newaxis, :, :]
    k_y = np.transpose(k_x, axes=[0, 2, 1])
    k_x = tf.convert_to_tensor(k_x, dtype=tf.complex64)
    k_x = tf.tile(k_x, multiples=(batch_size, 1, 1))
    k_y = tf.convert_to_tensor(k_y, dtype=tf.complex64)
    k_y = tf.tile(k_y, multiples=(batch_size, 1, 1))
    k_z_arg = tf.square(k) - (tf.square(k_x) + tf.square(k_y))
    k_z = tf.sqrt(k_z_arg)

    # Find shift amount
    theta_out = (
        np.pi
        / 180.0
        * tf.convert_to_tensor(
            np.tile(
                theta_sampling,
                np.size(wavelength_sampling) * np.size(phi_sampling),
            ),
            dtype=tf.float32,
        )
    )
    theta_out = theta_out[:, tf.newaxis, tf.newaxis]
    theta_out = tf.tile(theta_out, multiples=(1, n_pix, n_pix))
    theta = theta_out[:, 0, 0]
    theta = theta[:, np.newaxis, np.newaxis]
    y0_real = tf.tan(theta) * prop_dist

    phi = np.repeat(phi_sampling, np.size(theta_sampling))
    phi = np.tile(phi, np.size(wavelength_sampling))
    phi = np.pi / 180.0 * tf.convert_to_tensor(phi, dtype=tf.float32)
    phi = phi[:, tf.newaxis, tf.newaxis]
    phi = tf.tile(phi, multiples=(1, n_pix, n_pix))
    phi = phi[:, 0, 0]
    phi = phi[:, np.newaxis, np.newaxis]
    x0_real = tf.tan(phi) * prop_dist

    if lateral_shift is not None:
        x0_real = lateral_shift[0]
        y0_real = lateral_shift[1]

    y0 = tf.cast(y0_real, dtype=tf.complex64)
    x0 = tf.cast(x0_real, dtype=tf.complex64)

    propagator_arg = 1j * (k_z * prop_dist + k_x * x0 + k_y * y0)
    propagator = tf.exp(propagator_arg)

    # Limit transfer function bandwidth to prevent aliasing
    if use_antialiasing:
        S_x_src = n_pix * period
        S_y_src = S_x_src
        S_x_dest = n_pix * period
        S_y_dest = S_x_dest

        S_x_avg = 0.5 * (S_x_src + S_x_dest)
        kx_cond1 = S_x_avg < x0_real
        kx_cond2 = (-S_x_avg <= x0_real) & (x0_real < S_x_avg)

        S_y_avg = 0.5 * (S_y_src + S_y_dest)
        ky_cond1 = S_y_avg < y0_real
        ky_cond2 = (-S_y_avg <= y0_real) & (y0_real < S_y_avg)

        kx_limit_plus = k * tf.cast(
            1 / tf.math.sqrt((prop_dist / (x0_real + S_x_avg)) ** 2 + 1),
            dtype=tf.complex64,
        )
        kx_limit_minus = k * tf.cast(
            1 / tf.math.sqrt((prop_dist / (x0_real - S_x_avg)) ** 2 + 1),
            dtype=tf.complex64,
        )
        ky_limit_plus = k * tf.cast(
            1 / tf.math.sqrt((prop_dist / (y0_real + S_y_avg)) ** 2 + 1),
            dtype=tf.complex64,
        )
        ky_limit_minus = k * tf.cast(
            1 / tf.math.sqrt((prop_dist / (y0_real - S_y_avg)) ** 2 + 1),
            dtype=tf.complex64,
        )

        kx_region = tf.where(
            kx_cond1,
            _case1(k_x, k_y, k, kx_limit_minus, kx_limit_plus),
            _where_case2_case3(
                kx_cond2, k_x, k_y, k, kx_limit_minus, kx_limit_plus
            ),
        )

        ky_region = tf.where(
            ky_cond1,
            _case1(k_y, k_x, k, ky_limit_minus, ky_limit_plus),
            _where_case2_case3(
                ky_cond2, k_y, k_x, k, ky_limit_minus, ky_limit_plus
            ),
        )

        k_region = tf.cast(kx_region & ky_region, dtype=tf.complex64)

        return propagator * k_region
    else:
        return propagator

def _case1(k_x, k_y, k, kx_limit_minus, kx_limit_plus):
    """Case 1 of the transfer function bandwidth limiting."""
    kx_r1 = k_x.numpy() >= 0
    kx_r2 = (tf.square(k_x / kx_limit_minus) + tf.square(k_y / k)).numpy() >= 1
    kx_r3 = (tf.square(k_x / kx_limit_plus) + tf.square(k_y / k)).numpy() <= 1
    return kx_r1 & kx_r2 & kx_r3

def _case2(k_x, k_y, k, kx_limit_minus, kx_limit_plus):
    """Case 2 of the transfer function bandwidth limiting."""
    kx_r1 = k_x.numpy() <= 0
    kx_r2 = (tf.square(k_x / kx_limit_minus) + tf.square(k_y / k)).numpy() <= 1

    kx_r3 = k_x.numpy() > 0
    kx_r4 = (tf.square(k_x / kx_limit_plus) + tf.square(k_y / k)).numpy() <= 1
    return (kx_r1 & kx_r2) | (kx_r3 & kx_r4)

def _case3(k_x, k_y, k, kx_limit_minus, kx_limit_plus):
    """Case 3 of the transfer function bandwidth limiting."""
    kx_r1 = k_x.numpy() <= 0
    kx_r2 = (tf.square(k_x / kx_limit_plus) + tf.square(k_y / k)).numpy() >= 1
    kx_r3 = (tf.square(k_x / kx_limit_minus) + tf.square(k_y / k)).numpy() <= 1
    return kx_r1 & kx_r2 & kx_r3

def _where_case2_case3(kx_cond2, k_x, k_y, k, kx_limit_minus, kx_limit_plus):
    """Case 2 and 3 of the transfer function bandwidth limiting."""
    kx_region = tf.where(
        kx_cond2,
        _case2(k_x, k_y, k, kx_limit_minus, kx_limit_plus),
        _case3(k_x, k_y, k, kx_limit_minus, kx_limit_plus),
    )

    return kx_region

def propagate_with_propagator_batched(
    field: tf.Tensor,
    propagator: tf.Tensor,
    use_padding=True,
    upsampling=1,
) -> tf.Tensor:
    if use_padding:
        _, _, m = field.shape
        n = upsampling * m
        field = tf.transpose(field, perm=[1, 2, 0])
        field_real = tf.math.real(field)
        field_imag = tf.math.imag(field)
        field_real = tf.image.resize(field_real, [n, n], method="nearest")
        field_imag = tf.image.resize(field_imag, [n, n], method="nearest")
        field = tf.cast(field_real, dtype=tf.complex64) + 1j * tf.cast(
            field_imag, dtype=tf.complex64
        )
        field = tf.image.resize_with_crop_or_pad(field, 2 * n - 1, 2 * n - 1)
        field = tf.transpose(field, perm=[2, 0, 1])

    field_freq = tf.signal.fftshift(tf.signal.fft2d(field), axes=(1, 2))
    field_filtered = tf.signal.ifftshift(field_freq * propagator, axes=(1, 2))
    out = tf.signal.ifft2d(field_filtered)

    if use_padding:
        # Crop back down to n x n matrices
        out = tf.transpose(out, perm=[1, 2, 0])
        out = tf.image.resize_with_crop_or_pad(out, n, n)
        out = tf.transpose(out, perm=[2, 0, 1])

    return out

# def propagate_with_propagator(
#     field: tf.Tensor, propagator: tf.Tensor, use_padding=True, upsampling=1
# ) -> tf.Tensor:
#     """
#     Progragates the field through a distance using the transfer function.
#
#     Args:
#         field: The field to propagate.
#         propagator: The transfer function.
#         use_padding: Whether to use padding or not.
#         upsampling: The upsampling factor.
#     """
#     if use_padding:
#         _, m = field.shape
#         field = tf.transpose(field, perm=[1, 2, 0])
#         field_real = tf.math.real(field)
#         field_imag = tf.math.imag(field)
#         field_real = tf.image.resize(field_real, [m, m], method="nearest")
#         field_imag = tf.image.resize(field_imag, [m, m], method="nearest")
#         field = tf.cast(field_real, dtype=tf.complex64) + 1j * tf.cast(
#             field_imag, dtype=tf.complex64
#         )
#         field = tf.image.resize_with_crop_or_pad(field, 2 * m - 1, 2 * m - 1)
#         field = tf.transpose(field, perm=[2, 0, 1])
#
#     field_freq = tf.signal.fftshift(tf.signal.fft2d(field))
#     field_filtered = tf.signal.ifftshift(field_freq * propagator)
#     out = tf.signal.ifft2d(field_filtered)
#
#     if use_padding:
#         # Crop back down to m x m matrices
#         out = tf.image.resize_with_crop_or_pad(out, m, m)
#
#     return out

def get_incident_field_2d(
    field_props: FieldProperties,
) -> tf.Tensor:
    """Defines the input electric fields for the given wavelengths and field angles.

    Args:
        field_props: the field properties.

    Returns:
        Field2d: The incident field.
    """

    # Define the cartesian cross section
    # TODO: perioidicity for x and y seperately
    wavelength_base = field_props.wavelength
    theta_base = field_props.theta
    phi_base = field_props.phi

    # Define the cartesian cross section
    n_pix = field_props.n_pixels
    dx = field_props.period
    dy = field_props.period
    xa = np.linspace(0, n_pix - 1, n_pix) * dx  # x axis array
    xa = xa - np.mean(xa)  # center x axis at zero
    ya = np.linspace(0, n_pix - 1, n_pix) * dy  # y axis vector
    ya = ya - np.mean(ya)  # center y axis at zero
    [y_mesh, x_mesh] = np.meshgrid(ya, xa)
    x_mesh = x_mesh[np.newaxis, :, :]
    y_mesh = y_mesh[np.newaxis, :, :]

    lam0, theta, phi = utils.unravel_wavelength_theta_phi(
        wavelength_base, theta_base, phi_base
    )
    lam0 = lam0[:, tf.newaxis, tf.newaxis]
    theta = theta[:, tf.newaxis, tf.newaxis]
    phi = phi[:, tf.newaxis, tf.newaxis]

    phase_def = (
        2 * np.pi / lam0 * (np.sin(theta) * x_mesh + np.sin(phi) * y_mesh)
    )
    phase_def = tf.cast(phase_def, dtype=tf.complex64)

    tensor = tf.exp(1j * phase_def)
    total_energy = tf.reduce_sum(tf.math.abs(tensor) ** 2)
    tensor /= tf.sqrt(tf.cast(total_energy, dtype=tf.complex64))

    return Field2D(
        tensor=tensor,
        n_pixels=field_props.n_pixels,
        wavelength=field_props.wavelength,
        theta=field_props.theta,
        phi=field_props.phi,
        period=field_props.period,
        upsampling=field_props.upsampling,
        use_padding=field_props.use_padding,
        use_antialiasing=field_props.use_antialiasing,
    )


import os.path as osp
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch
import util

class Net(torch.nn.Module):
    def __init__(self,device, in_channels,K, nb_chev_filter, num_time_filter, T, P, num_nodes, time_strides,cheb_polynomials):
        super(Net, self).__init__()
        self.BlockList1 = GAT_block(device = device, in_channels = in_channels, K = K, nb_chev_filter = nb_chev_filter, num_time_filter =  num_time_filter,T =  T, num_nodes = num_nodes, time_strides = time_strides, cheb_polynomials = cheb_polynomials)
        self.BlockList2 = GAT_block(device = device, in_channels = 64, K = K, nb_chev_filter = nb_chev_filter, num_time_filter =  num_time_filter, T =  T, num_nodes = num_nodes, time_strides = time_strides, cheb_polynomials = cheb_polynomials)

        self.final_conv = nn.Conv2d(int(T / time_strides), P, kernel_size=(1, num_time_filter))

    def forward(self, x):
        '''
        :param x: (Batch_size, Num_nodes, F_in, T_in)
        :return: (B, Num_nodes, T_out)
        '''
        list = torch.split(x, 1, dim=3)

        x = list[0].permute(0,2,3,1)

        #x = x.permute(0,2,3,1)
        x = self.BlockList1(x)
        x = self.BlockList2(x)
        x = self.final_conv(x.permute(0, 3, 1, 2))[:, :, :, -1].permute(0, 1, 2)

        return x

class GAT_block(nn.Module):
    def __init__(self,device, in_channels,K,nb_chev_filter, num_time_filter, T, num_nodes, time_strides,cheb_polynomials):
        super(GAT_block, self).__init__()
        self.TAL = Temporal_Attention_Layer(device,in_channels, num_nodes,T)
        self.SAL = Spatial_Attention_Layer(device,in_channels, num_nodes,T)
        self.cheb_conv_SAt = cheb_conv_withSAt(K, cheb_polynomials, in_channels, nb_chev_filter)
        self.time_conv = nn.Conv2d(nb_chev_filter, num_time_filter, kernel_size=(1, 3), stride=(1, time_strides), padding=(0, 1))
        self.residual_conv = nn.Conv2d(in_channels, num_time_filter, kernel_size=(1, 1), stride=(1, time_strides))
        self.ln = nn.LayerNorm(num_time_filter)

    def forward(self,x):
        '''
        :param x: (Batch_size, Num_nodes, F_in, T)
        :return: (Batch_size, Num_nodes, nb_time_filter, T)
        '''

        batch_size, num_nodes, num_features, num_timesteps = x.shape

        # temporal attention
        temporal_attention = self.TAL(x) # (Batchsize,T, T)
        x_TAt = torch.matmul(x.reshape(batch_size, -1, num_timesteps), temporal_attention).reshape(batch_size,
                                                                                num_nodes,
                                                                                num_features,
                                                                                num_timesteps)
        spatial_attention = self.SAL(x_TAt) # (Batchsize,Num_nodes,Num_features,T)

        # cheb gcn
        spatial_gcn = self.cheb_conv_SAt(x, spatial_attention)  # (b,N,F,T)
        # spatial_gcn = self.cheb_conv(x)

        # convolution along the time axis
        time_conv_output = self.time_conv(spatial_gcn.permute(0, 2, 1, 3))  # (b,N,F,T)->(b,F,N,T) 用(1,3)的卷积核去做->(b,F,N,T)

        # residual shortcut
        x_residual = self.residual_conv(x.permute(0, 2, 1, 3))  # (b,N,F,T)->(b,F,N,T) 用(1,1)的卷积核去做->(b,F,N,T)

        x_residual = self.ln(F.relu(x_residual + time_conv_output).permute(0, 3, 2, 1)).permute(0, 2, 3, 1)
        # (b,F,N,T)->(b,T,N,F) -ln-> (b,T,N,F)->(b,N,F,T)

        return x_residual

class Temporal_Attention_Layer(nn.Module):
    def __init__(self, DEVICE, in_channels, num_of_vertices, num_of_timesteps):
        super(Temporal_Attention_Layer, self).__init__()
        self.U1 = nn.Parameter(torch.FloatTensor(num_of_vertices).to(DEVICE))
        self.U2 = nn.Parameter(torch.FloatTensor(in_channels, num_of_vertices).to(DEVICE))
        self.U3 = nn.Parameter(torch.FloatTensor(in_channels).to(DEVICE))
        self.be = nn.Parameter(torch.FloatTensor(1, num_of_timesteps, num_of_timesteps).to(DEVICE))
        self.Ve = nn.Parameter(torch.FloatTensor(num_of_timesteps, num_of_timesteps).to(DEVICE))

    def forward(self, x):
        '''
        :param x: (Batch_size, N, F_in, T)
        :return: (B, T, T)
        '''
        batch_size, num_nodes, num_features, num_timesteps = x.shape
        a = torch.matmul(x.permute(0, 3, 2, 1), self.U1)
        lhs = torch.matmul(a, self.U2)
        # x:(B, N, F_in, T) -> (B, T, F_in, N)
        # (B, T, F_in, N)(N) -> (B,T,F_in)
        # (B,T,F_in)(F_in,N)->(B,T,N)

        rhs = torch.matmul(self.U3, x)  # (F)(B,N,F,T)->(B, N, T)

        product = torch.matmul(lhs, rhs)  # (B,T,N)(B,N,T)->(B,T,T)

        E = torch.matmul(self.Ve, torch.sigmoid(product + self.be))  # (B, T, T)

        E_normalized = F.softmax(E, dim=1)

        return E_normalized

class Spatial_Attention_Layer(nn.Module):
    def __init__(self, DEVICE, in_channels, num_of_vertices, num_of_timesteps):
        super(Spatial_Attention_Layer, self).__init__()
        self.W1 = nn.Parameter(torch.FloatTensor(num_of_timesteps).to(DEVICE))
        self.W2 = nn.Parameter(torch.FloatTensor(in_channels, num_of_timesteps).to(DEVICE))
        self.W3 = nn.Parameter(torch.FloatTensor(in_channels).to(DEVICE))
        self.bs = nn.Parameter(torch.FloatTensor(1, num_of_vertices, num_of_vertices).to(DEVICE))
        self.Vs = nn.Parameter(torch.FloatTensor(num_of_vertices, num_of_vertices).to(DEVICE))
    def forward(self, x):
        '''
        :param x: (batch_size, N, F_in, T)
        :return: (B,N,N)
        '''

        lhs = torch.matmul(torch.matmul(x, self.W1), self.W2)  # (b,N,F,T)(T)->(b,N,F)(F,T)->(b,N,T)

        rhs = torch.matmul(self.W3, x).transpose(-1, -2)  # (F)(b,N,F,T)->(b,N,T)->(b,T,N)

        product = torch.matmul(lhs, rhs)  # (b,N,T)(b,T,N) -> (B, N, N)

        S = torch.matmul(self.Vs, torch.sigmoid(product + self.bs))  # (N,N)(B, N, N)->(B,N,N)

        S_normalized = F.softmax(S, dim=1)

        return S_normalized

class cheb_conv_withSAt(nn.Module):
    '''
    K-order chebyshev graph convolution
    '''

    def __init__(self, K, cheb_polynomials, in_channels, out_channels):
        '''
        :param K: int
        :param in_channles: int, num of channels in the input sequence
        :param out_channels: int, num of channels in the output sequence
        '''
        super(cheb_conv_withSAt, self).__init__()
        self.K = K
        self.cheb_polynomials = cheb_polynomials
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.DEVICE = cheb_polynomials[0].device
        self.Theta = nn.ParameterList([nn.Parameter(torch.FloatTensor(in_channels, out_channels).to(self.DEVICE)) for _ in range(K)])

    def forward(self, x, spatial_attention):
        '''
        Chebyshev graph convolution operation
        :param x: (batch_size, N, F_in, T)
        :return: (batch_size, N, F_out, T)
        '''

        batch_size, num_of_vertices, in_channels, num_of_timesteps = x.shape

        outputs = []

        for time_step in range(num_of_timesteps):

            graph_signal = x[:, :, :, time_step]  # (b, N, F_in)

            output = torch.zeros(batch_size, num_of_vertices, self.out_channels).to(self.DEVICE)  # (b, N, F_out)

            for k in range(self.K):

                T_k = self.cheb_polynomials[k]  # (N,N)

                T_k_with_at = T_k.mul(spatial_attention)   # (N,N)*(N,N) = (N,N) 多行和为1, 按着列进行归一化

                theta_k = self.Theta[k]  # (in_channel, out_channel)

                rhs = T_k_with_at.permute(0, 2, 1).matmul(graph_signal)  # (N, N)(b, N, F_in) = (b, N, F_in) 因为是左乘，所以多行和为1变为多列和为1，即一行之和为1，进行左乘

                output = output + rhs.matmul(theta_k)  # (b, N, F_in)(F_in, F_out) = (b, N, F_out)

            outputs.append(output.unsqueeze(-1))  # (b, N, F_out, 1)

        return F.relu(torch.cat(outputs, dim=-1))  # (b, N, F_out, T)

def model(args, A):
    '''
    :param args.device
    :param args.in_channels
    :param K:
    :param args.nb_block
    :param nb_time_filter:
    :param args.S:
    :param args.P:
    :param args.num_nodes:
    :param time_strides
    :param edge_index
    :param edge_weight
    :return:
    '''
    K = 3
    nb_chev_filter = 64
    L_tilde = util.scaled_Laplacian(A)
    cheb_polynomials = [torch.from_numpy(i).type(torch.FloatTensor).to(args.device) for i in util.cheb_polynomial(L_tilde, K)]
    model = Net(args.device, args.in_channels,K, nb_chev_filter, 64, args.T, args.P, args.num_nodes, 1, cheb_polynomials)

    for p in model.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)
        else:
            nn.init.uniform_(p)

    return model


import telebot
import random
import logging
from telebot import types

list1 = ["Пизда", "Хуй на"]

logger = telebot.logger
telebot.logger.setLevel(logging.DEBUG)

bot = telebot.TeleBot("5345165100:AAF04G_rkh_uoxECYrTaOw_nTwV8PczpfX4", parse_mode=None)

@bot.message_handler(content_types=['text'])
def send_text(message):
    with open('tq.txt') as f:
        lines = f.readlines()
    try:
        with open("tq.txt", "a+") as file_object:
            print(message.text)
            file_object.seek(0)
            data = file_object.read(100)
            if len(data) > 0 :
                file_object.write("\n")
            file_object.write(message.text)
    except:
        pass
    if message.text.lower() == "да" or message.text.lower() == "@mrgshit_bot да":
        bot.reply_to(message, random.choice(list1))
    elif "твоя" in message.text.lower() and "мама" in message.text.lower():
        bot.reply_to(message, "Польская кошкодевочка с миской риса")
    elif "твой" in message.text.lower() and "папа" in message.text.lower():
        bot.reply_to(message, "Гений, миллиардер, плейбой, филантроп")
    elif message.text.lower() == "асу" or message.text.lower() == "@mrgshit_bot асу":
        bot.reply_to(message, "Пов: ввели санкции на буквы Ж Д А Ю")
    elif message.text.lower() == 'пизда':
        bot.reply_to(message, "Это моя фраза, чмо")
    elif message.text.lower() == 'хуй на':
        bot.reply_to(message, "Соси два")
    elif "алина" in message.text.lower():
        bot.reply_to(message, "Алина шлюха")
    elif "путин" in message.text.lower():
        bot.reply_to(message, "Путин хуйло")
    else:
        if random.randint(0,100) < 15 or "@mrgshit_bot" in message.text.lower():
            try:
                bot.reply_to(message, random.choice(lines))
            except:
                pass

@bot.message_handler(content_types=['document', 'audio'])
def handle_docs_audio(message):
    bot.reply_to(message, "2")
bot.polling()


# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import os
import re

from setuptools import find_packages, setup

THIS_DIR = os.path.dirname(os.path.realpath(__name__))

def read(*parts):
    with open(os.path.join(THIS_DIR, *parts)) as f:
        return f.read()

def get_version():
    return re.findall(
        '__version__ = "([\d\.]+)"', read("firefox_ui_harness", "__init__.py"), re.M
    )[0]

long_description = """Custom Marionette runner classes and entry scripts for Firefox Desktop
specific Marionette tests.
"""

setup(
    name="firefox-ui-harness",
    version=get_version(),
    description="Firefox UI Harness",
    long_description=long_description,
    classifiers=[
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.5",
        "License :: OSI Approved :: Mozilla Public License 2.0 (MPL 2.0)",
    ],
    # Get strings from http://pypi.python.org/pypi?%3Aaction=list_classifiers
    keywords="mozilla",
    author="DevTools",
    author_email="dev-webdriver@mozilla.org",
    license="MPL",
    packages=find_packages(),
    include_package_data=True,
    zip_safe=False,
    install_requires=read("requirements.txt").splitlines(),
    entry_points="""
        [console_scripts]
        firefox-ui-functional = firefox_ui_harness.cli_functional:cli
      """,
)


import streamlit as st
import pandas as pd
import plotly.express as px
from email.policy import default
from operator import index
import matplotlib.pyplot as plt
from streamlit_option_menu import option_menu

import numpy as np
from flask import Flask, request, jsonify, render_template, url_for
import joblib
from sklearn import svm

@st.cache(allow_output_mutation=True, persist = True)
def cargar_de_datos():
    leer_columnas = ['ESTU_DEPTO_RESIDE','ESTU_GENERO','FAMI_TIENECOMPUTADOR','FAMI_TIENEINTERNET','FAMI_ESTRATOVIVIENDA','FAMI_COMECARNEPESCADOHUEVO','PUNT_GLOBAL']
    df = pd.read_csv("icfes_data.csv",low_memory=False, usecols=leer_columnas)
    return df

st.set_page_config(
     page_title="PREDICCIÓN ICFES",
     layout="wide",
     initial_sidebar_state="expanded"
 )

df = cargar_de_datos().copy()
with st.sidebar:
    st.image('logo.png')
    menu = option_menu(
        menu_title = "BIENVENIDO",
        options = ['Inicio','Exploración', 'Predicción','Contactenos'],
        default_index =0
    )

if menu == 'Inicio':

    # Título
    html_temp = """

ICFES COLOMBIA 2018-2021

    """
    st.markdown(html_temp,unsafe_allow_html=True)

    st.subheader("El dataset en el cual estamos trabajando es *Icfes Colombia 2018-2021* en este encontramos datos relacionados con las características de estudiantes que realizaron las  pruebas icfes en unos periodos determinados por año. En el dataset original se describe información   general, académica, familiar y socioeconómica del los estudiantes.")
    ver = st.button('Mostrar datos')
    if ver:
        st.dataframe(df)

        st.text("PUNTAJE GLOBAL POR DEPARTAMENTOS")
        cols = [i for i in df.columns if i.startswith('PUNT_')]
        inf = df.groupby(["ESTU_DEPTO_RESIDE"])[cols].mean()
        st.dataframe(inf)

        st.text("PUNTAJE GLOBAL - INTERNET")
        internet= df.groupby(["FAMI_TIENEINTERNET"])[["PUNT_GLOBAL"]].mean()
        st.dataframe(internet)

        st.text("PUNTAJE GLOBAL - COMPUTADOR")
        internet= df.groupby(["FAMI_TIENECOMPUTADOR"])[["PUNT_GLOBAL"]].mean()
        st.dataframe(internet)

        st.text("PUNTAJE GLOBAL POR ESTRATO")
        internet= df.groupby(["FAMI_ESTRATOVIVIENDA"])[["PUNT_GLOBAL"]].mean()
        st.dataframe(internet)

        st.text("PUNTAJE GLOBAL - ALIMENTACIÓN DEL ESTUDIANTE(CARNES-PESCADO-HUEVOS)")
        internet= df.groupby(["FAMI_COMECARNEPESCADOHUEVO"])[["PUNT_GLOBAL"]].mean()
        st.dataframe(internet)

elif menu == 'Exploración':

    # Título
    html_temp = """

EXPLORACIÓN

    """
    st.markdown(html_temp,unsafe_allow_html=True)

    a1 = df.groupby(by=['ESTU_DEPTO_RESIDE']).mean()[['PUNT_GLOBAL']]

    fig1 = px.bar(

        a1,
        x = "PUNT_GLOBAL",
        y = a1.index,
        orientation = "h",
        height=700, width = 1200,
        title= " PUNTAJE GLOBAL POR DEPARTAMENTOS ",
        color_discrete_sequence = ['#618AE8'] * len(a1),
        template = "plotly_white"
        ).update_layout(xaxis_title="PRORMEDIO DE PUNTAJE GLOBAL", yaxis_title="DEPARTAMENTOS",yaxis={'categoryorder':'total ascending'})

    fig1

    a2 = df.groupby(by=['FAMI_ESTRATOVIVIENDA']).mean()[['PUNT_GLOBAL']]

    fig2 = px.bar(

        a2,
        x = "PUNT_GLOBAL",
        y = a2.index,
        height=400, width = 1200,
        title= " PROMEDIO DE PUNTAJE GLOBAL POR ESTRATOS ",
        color_discrete_sequence = ['#618AE8'] * len(a2),
        template = "plotly_white"
        ).update_layout(xaxis_title=" PUNTAJE DE PUNTAJE GLOBAL", yaxis_title="ESTRATO",yaxis={'categoryorder':'total ascending'})

    fig2

    a3 = df.groupby(by=['FAMI_COMECARNEPESCADOHUEVO']).mean()[['PUNT_GLOBAL']]
    fig3 = px.bar(

        a3,
        x = "PUNT_GLOBAL",
        y = a3.index,
        height=400, width = 1200,
        title= " PUNTAJE GLOBAL POR ALIMENTACIÓN ",
        color_discrete_sequence = ['#618AE8'] * len(a3),
        template = "plotly_white"
        ).update_layout(xaxis_title="PROMEDIO DE PUNTAJE GLOBAL", yaxis_title="VECES QUE CONSUME CARNES-PESCADO-HUEVO",yaxis={'categoryorder':'total ascending'})

    fig3

    a3 = df.groupby(by=['ESTU_GENERO']).mean()[['PUNT_GLOBAL']]

    fig4 = px.bar(

        a3,
        x = "PUNT_GLOBAL",
        y = a3.index,
        orientation = "h",
        height=700, width = 1200,
        title= " PUNTAJE GLOBAL POR GENERO ",
        color_discrete_sequence = ['#618AE8'] * len(a3),
        template = "plotly_white"
        ).update_layout(xaxis_title="PRORMEDIO DE PUNTAJE GLOBAL", yaxis_title="GENERO",yaxis={'categoryorder':'total ascending'})

    fig4

    def genero_est():
            genero= df['FAMI_TIENEINTERNET'].value_counts()

            datos = ['Si', 'No']
            colores = ["#618AE8","#5c60e6"]

            plt.pie(genero, labels=datos,autopct="%0.lf %%", colors= colores,radius=0.5)

            plt.title(' PORCENTAJE CON INTERNET/SIN INTERNET')
            plt.show()

            st.pyplot(plt)

    genero_est()

#PREDICCION

elif menu == 'Predicción':

    # Path del modelo preentrenado
    MODEL_PATH = 'my_model.pkl'

    # Se recibe la imagen y el modelo, devuelve la predicción
    def model_prediction(x_in, model):

        x = np.asarray(x_in).reshape(1,-1)
        preds=model.predict(x)

        return preds

    with open(MODEL_PATH, 'rb') as file:
        model = joblib.load(file)

    # Título
    html_temp = """

SISTEMA DE PREDICCIÓN APROBACIÓN PRUEBAS ICFES EN COLOMBIA

    """
    st.markdown(html_temp,unsafe_allow_html=True)

    st.subheader("Se pretende predecir si el estudiante gana o pierde las pruebas icfes teniendo en cuenta las variables socioeconomicas ")

    st.subheader(" Nota: Gana las pruebas icfes apartir de un puntacion mayor o igual a 250, de lo contrario pierde la prueba")

    estratos = {'Sin Estrato':0,'Estrato 1':1 , 'Estrato 2':2 ,'Estrato 3':3 , 'Estrato 4':4 ,'Estrato 5':5,'Estrato 6':6}

    genero = {'Femenino':0,'Masculino':1}

    internet = {'No':0,'Si':1}
    computador = {'No':0,'Si':1}

    N = st.selectbox("Genero" , ('Masculino','Femenino'))
    C = st.selectbox("Tiene Computador:", ('Si','No'))
    K = st.selectbox("Tiene Internet:", ('Si','No'))
    P = st.selectbox("Estrato de Vivienda:", ('Sin Estrato','Estrato 1','Estrato 2','Estrato 3','Estrato 4','Estrato 5','Estrato 6'))

    N = genero[N]
    P = estratos[P]
    K = internet[K]
    C = computador[C]

    # El botón predicción se usa para iniciar el procesamiento
    if st.button("Predicción"):
        #x_in = list(np.float_((Datos.title().split('\t'))))
        x_in =[np.float_(N),
                    np.float_(P),
                    np.float_(K),
                    np.float_(C)
            ]
        predictS = model_prediction(x_in, model)

        pred = format(predictS[0]).upper()

        if(int(pred) == 1):
            a = "PERDIO"
        else:
            a = "GANO"

        st.success('Se predice que el estudiante ' + a + " las pruebas ICFES")

elif menu == 'Contactenos':

        # Título
    html_temp = """

GRUPO 1

JORGE LUIS MONTES GOMEZ

CUADRADO TRUJILLO DIEGO.

PERALTA OSORIO LUIS

DANIELA TORRES MINA

DIAZ MONTIEL JEAN CARLOS

MACEA PIMIENTA EDER

    """
    st.markdown(html_temp,unsafe_allow_html=True)


from setuptools import setup, find_packages

requires = [
    'click',
    'biopython>=1.7',
    'pandas'
    ]

setup(
    name="tnseq2",
    version="1.0",
    author="Anna Sintsova",
    author_email="ansintsova@ethz.ch",
    description=("Bioinformatic toolkit for barcode mapping and counting"),
    license="LICENSE",
    keywords="tnseq",
    install_requires=requires,
    #url = "https://github.com/SushiLab/TNSEQ_DEV",
    packages=find_packages(),
    entry_points={
        'console_scripts': ['tnseq2=tnseq2.main:main'],
    }
)


import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pandas_datareader as pdr
import yfinance as yf
from pykalman import KalmanFilter

def plot_spread(prices, errors, stds, neg_stds):
    # Plot the spread
    plt.figure(figsize=(20, 6))
    plt.plot(prices.index[1:], errors, color='blue', label='KF Spread')
    plt.plot(prices.index[1:], stds, linestyle='--', color='red', label='Std Dev')
    plt.plot(prices.index[1:], neg_stds, linestyle='--', color='red', label='-Std Dev')
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.xlabel("Days")
    plt.ylabel("KF Spread")
    plt.title("Kalman Filter Spread")
    plt.legend()
    plt.show()

def plot_spread_rollingOLS(prices, errors, stds, neg_stds, window):
    # Plot the spread
    plt.figure(figsize=(20, 6))
    plt.plot(prices.index[2*window-2:], errors, color='blue', label='z-Score')
    plt.plot(prices.index[2*window-1:], stds, linestyle='--', color='red', label='+1')
    plt.plot(prices.index[2*window-1:], neg_stds, linestyle='--', color='red', label='-1')
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.xlabel("Days")
    plt.ylabel("z-Score")
    plt.title("z-Score of Spread")
    plt.legend()
    plt.show()

def plot_nav(prices, nav, window):
    # Plot Unit NAV
    plt.figure(figsize=(20, 10))
    plt.plot(prices.index[window-1:], nav, color='blue', label='Cumulative Unit NAV')
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.xlabel('Days')
    plt.ylabel('Unit NAV')
    plt.title('Cumulative Unit NAV')
    plt.legend()
    plt.show()

def plot_nav_sp500(prices, nav, sp500_nav, USTreasury_nav):
    # Get S&P 500 data from Yahoo Finance
    # index = prices.index.strftime('%Y-%m-%d')
    # sp500 = yf.download('^GSPC', start=index[0], end=index[-1])['Adj Close']
    # sp500 = pdr.get_data_yahoo('^GSPC', start=prices.index[0], end=prices.index[-1])['Adj Close']

    # Plot Unit NAV and S&P 500
    fig, ax = plt.subplots(figsize=(20, 10))
    ax.plot(prices.index, nav, color='red', label='Portfolio')
    ax.plot(prices.index[1:], sp500_nav, color='blue', label='S&P 500')
    ax.plot(prices.index[1:], USTreasury_nav, color='lightblue', label='10 Year US Treasury')
    ax.grid(True, linestyle='--', alpha=0.5)
    ax.set_xlabel('Days')
    ax.set_ylabel('Cumulative Unit NAV')
    ax.set_title('Cumulative Unit NAV')
    ax.legend()
    plt.show()

def plot_nav_sp500_rollingOLS(prices, nav, sp500_nav, USTreasury_nav, window):
    # Get S&P 500 data from Yahoo Finance
    # index = prices.index.strftime('%Y-%m-%d')
    # sp500 = yf.download('^GSPC', start=index[0], end=index[-1])['Adj Close']
    # sp500 = pdr.get_data_yahoo('^GSPC', start=prices.index[0], end=prices.index[-1])['Adj Close']

    # Plot Unit NAV and S&P 500
    fig, ax = plt.subplots(figsize=(20, 10))
    ax.plot(prices.index[2*window-2:], nav, color='red', label='Portfolio')
    ax.plot(prices.index[2*window-1:], sp500_nav, color='blue', label='S&P 500')
    ax.plot(prices.index[2*window-1:], USTreasury_nav, color='lightblue', label='10 Year US Treasury')
    ax.grid(True, linestyle='--', alpha=0.5)
    ax.set_xlabel('Days')
    ax.set_ylabel('Cumulative Unit NAV')
    ax.set_title('Cumulative Unit NAV')
    ax.legend()
    plt.show()

def plot_cum_returns(prices, cum_returns):
    # Plot cumulative returns
    plt.figure(figsize=(20, 10))
    plt.plot(prices.index[1:], cum_returns, color='blue', label='Cumulative Returns')
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.xlabel('Days')
    plt.ylabel('Cumulative returns')
    plt.title('Cumulative Returns')
    plt.legend()
    plt.show()

def plot_drawdown(prices, drawdowns):
    # Plot the drawdowns
    fig, ax = plt.subplots(figsize=(20, 10))
    ax.plot(prices.index, drawdowns, color='blue')
    # ax.plot(prices.index[1:], drawdowns, color='blue')
    ax.fill_between(prices.index, drawdowns, 0, color='red', alpha=0.3)
    # ax.fill_between(prices.index[1:], drawdowns, 0, color='red', alpha=0.3)
    ax.set_xlabel('Time')
    ax.set_ylabel('Drawdown')
    ax.set_title('Drawdown')

    plt.show()

def plot_drawdown_sp500(prices, drawdowns, sp500_drawdowns, USTreasury_drawdowns):
    # Get S&P 500 data from Yahoo Finance
    # index = prices.index.strftime('%Y-%m-%d')
    # sp500 = yf.download('^GSPC', start=index[0], end=index[-1])['Adj Close']
    # sp500_nav = sp500 / sp500.iloc[0]
    # sp500_highwatermarks = np.maximum.accumulate(sp500_nav)
    # sp500_drawdowns = (sp500_nav - sp500_highwatermarks) / sp500_highwatermarks
    # Plot the drawdowns
    fig, ax = plt.subplots(figsize=(20, 10))
    ax.plot(prices.index, drawdowns, color='red', label='Portfolio Drawdown')
    # ax.plot(prices.index[1:], drawdowns, color='blue')
    ax.fill_between(prices.index, drawdowns, 0, color='red', alpha=0.3)
    # ax.fill_between(prices.index[1:], drawdowns, 0, color='red', alpha=0.3)

    ax.plot(prices.index[1:], sp500_drawdowns, color='blue', label='S&P 500 Drawdown')
    ax.fill_between(prices.index[1:], sp500_drawdowns, 0, color='blue', alpha=0.3)

    ax.plot(prices.index[1:], USTreasury_drawdowns, color='orange', label='10 Year US Treasury Drawdown')
    ax.fill_between(prices.index[1:], USTreasury_drawdowns, 0, color='orange', alpha=0.3)
    ax.grid(True, linestyle='--', alpha=0.5)
    ax.set_xlabel('Time')
    ax.set_ylabel('Drawdown')
    ax.set_title('Drawdown')
    ax.legend()
    plt.show()

def plot_drawdown_sp500_rollingOLS(prices, drawdowns, sp500_drawdowns, USTreasury_drawdowns, window):
    # Get S&P 500 data from Yahoo Finance
    # index = prices.index.strftime('%Y-%m-%d')
    # sp500 = yf.download('^GSPC', start=index[0], end=index[-1])['Adj Close']
    # sp500_nav = sp500 / sp500.iloc[0]
    # sp500_highwatermarks = np.maximum.accumulate(sp500_nav)
    # sp500_drawdowns = (sp500_nav - sp500_highwatermarks) / sp500_highwatermarks
    # Plot the drawdowns
    fig, ax = plt.subplots(figsize=(20, 10))
    ax.plot(prices.index[2*window-2:], drawdowns, color='red', label='Portfolio Drawdown')
    # ax.plot(prices.index[1:], drawdowns, color='blue')
    ax.fill_between(prices.index[2*window-2:], drawdowns, 0, color='red', alpha=0.3)
    # ax.fill_between(prices.index[1:], drawdowns, 0, color='red', alpha=0.3)

    ax.plot(prices.index[2*window-1:], sp500_drawdowns, color='blue', label='S&P 500 Drawdown')
    ax.fill_between(prices.index[2*window-1:], sp500_drawdowns, 0, color='blue', alpha=0.3)

    ax.plot(prices.index[2*window-1:], USTreasury_drawdowns, color='orange', label='10 Year US Treasury Drawdown')
    ax.fill_between(prices.index[2*window-1:], USTreasury_drawdowns, 0, color='orange', alpha=0.3)
    ax.grid(True, linestyle='--', alpha=0.5)
    ax.set_xlabel('Time')
    ax.set_ylabel('Drawdown')
    ax.set_title('Drawdown')
    ax.legend()
    plt.show()

def calculate_daily_returns(portfolio_nav):
    daily_returns = []

    for i in range(1, len(portfolio_nav)):
        daily_return = (portfolio_nav[i] - portfolio_nav[i - 1]) / portfolio_nav[i - 1]
        daily_returns.append(daily_return)

    return daily_returns

def calculate_drawdown(portfolio_nav):
    max_so_far = 0
    drawdowns = []

    for nav in portfolio_nav:
        max_so_far = max(max_so_far, nav)
        drawdown = -(max_so_far - nav) / max_so_far
        drawdowns.append(drawdown)

    return drawdowns

def max_drawdown(daily_returns):
    # Convert daily returns to a Pandas Series object
    daily_returns = pd.Series(daily_returns)

    # Compute the cumulative returns at each time point
    cumulative_unit_nav = (1 + daily_returns).cumprod()
    # cumulative_returns = np.cumprod(1 + daily_returns) - 1

    # Compute the highwatermarks at each time point
    highwatermarks = np.maximum.accumulate(cumulative_unit_nav)

    # Compute the drawdowns at each time point
    drawdowns = (cumulative_unit_nav - highwatermarks) / highwatermarks

    # Compute the maximum drawdown and the start and end dates of the maximum drawdown period
    max_dd = np.min(drawdowns)
    # max_dd_start_date = np.argmax(np.maximum.accumulate(cumulative_unit_nav) - cumulative_unit_nav[:np.argmax(np.maximum.accumulate(cumulative_unit_nav))])
    # max_dd_end_date = np.argmax(cumulative_unit_nav[:np.argmax(np.maximum.accumulate(cumulative_unit_nav))] + max_dd)

    return max_dd

def pnl_metrics(daily_returns, sp500_nav, sp500_drawdowns, USTreasury_nav, USTreasury_drawdowns):
    # Calculate annualized return
    annualized_return = np.mean(daily_returns) * 252

    # Calculate annualized volatility
    annualized_volatility = np.std(daily_returns) * np.sqrt(252)

    # Calculate Sharpe Ratio
    sharpe_ratio = (annualized_return - 0.02) / annualized_volatility

    # Calculate Max Drawdown
    max_dd = max_drawdown(daily_returns)
    # max_drawdown = max_drawdown(cum_returns)

    # Calculate S&P 500 annualized return, volatility and max drawdown
    sp500_annualized_return = np.mean(sp500_nav.pct_change().dropna()) * 252
    sp500_annualized_volatility = np.std(sp500_nav.pct_change().dropna()) * np.sqrt(252)
    sp500_sharpe_ratio = (sp500_annualized_return - 0.02) / sp500_annualized_volatility
    sp500_max_dd = np.min(sp500_drawdowns)

    # Calculate US Treasury annualized return, volatility and max drawdown
    USTreasury_annualized_return = np.mean(USTreasury_nav.pct_change().dropna()) * 252
    USTreasury_annualized_volatility = np.std(USTreasury_nav.pct_change().dropna()) * np.sqrt(252)
    USTreasury_max_dd = np.min(USTreasury_drawdowns)

    results = pd.DataFrame({'KF Pairs Trading': [annualized_return, annualized_volatility, sharpe_ratio, max_dd],
                           'S&P 500': [sp500_annualized_return, sp500_annualized_volatility, sp500_sharpe_ratio, sp500_max_dd],
                           '10 Year US Treasury': [USTreasury_annualized_return, USTreasury_annualized_volatility, np.nan, USTreasury_max_dd]},
                           index=['Annualized Return', 'Annualized Volatility', 'Sharpe Ratio', 'Max Drawdown'])
    # results = results.to_string(header=True,
    #                             col_space=10,
    #                             justify='left',
    #                             float_format="%.3f")
    return results


from torch.utils.data import Dataset
import numpy as np
import os
from datasets.preprocess import *
from datasets.data_io import *
from imageio import imread, imsave, imwrite

# all of existed dataset preprocessed by Jin Liu (only for training and validation)
class MVSDataset(Dataset):
    def __init__(self, data_folder, set_name, mode, view_num, interval_scale):
        super(MVSDataset, self).__init__()
        self.data_folder = data_folder
        self.set_name = set_name
        self.mode = mode
        self.view_num = view_num
        self.interval_scale = interval_scale
        self.counter = 0

        assert self.mode in ["train", "test", "val"]
        self.sample_list = self.build_list()
        self.sample_num = len(self.sample_list)

    def build_list(self):
        # Prepare all training samples
        """ generate data paths for whu dataset """

        data_folder = self.data_folder

        if self.set_name == 'dtu':
            sample_list = self.dtu_list(data_folder, self.set_name, gt_fext='.pfm')
        elif self.set_name == 'BlendedMVS':
            sample_list = self.BlendedMVS_list(data_folder, self.set_name, gt_fext='.pfm')
        elif self.set_name == 'whu_omvs':
            sample_list = self.ObliqueWhu_list(data_folder, self.set_name, gt_fext='.exr')
        elif self.set_name == 'whu_mvs':
            sample_list = self.whu_list(data_folder, self.set_name, gt_fext='.png')
        else:
            raise Exception("{}? Not implemented yet!".format(self.set_name))

        return sample_list

    def whu_list(self, data_folder, sat_name, gt_fext='.png'):
        sample_list = []

        # image index
        train_cluster_path = data_folder + '/index.txt'
        data_cluster = open(train_cluster_path).read().split()

        # pair
        view_pair_path = data_folder + '/pair.txt'
        ref_indexs = []
        src_indexs = []
        with open(view_pair_path) as f:
            cluster_num = int(f.readline().rstrip())
            for idx in range(cluster_num):
                ref_index = int(f.readline().rstrip())
                src_index = [int(x) for x in f.readline().rstrip().split()][1:]
                ref_indexs.append(ref_index)
                src_indexs.append(src_index)

        # for each data scene
        for i in data_cluster:
            image_folder = os.path.join(data_folder, ('Images/%s' % i)).replace("\\", "/")
            cam_folder = os.path.join(data_folder, ('Cams/%s' % i)).replace("\\", "/")
            depth_folder = os.path.join(data_folder, ('Depths/%s' % i)).replace("\\", "/")

            # for each view
            for ref_ind, view_inds in zip(ref_indexs, src_indexs):  # 0-4
                image_folder2 = os.path.join(image_folder, ('%d' % ref_ind)).replace("\\", "/")
                image_files = sorted(os.listdir(image_folder2))

                view_cnts = min(self.view_num, len(view_inds) + 1)

                for j in range(0, int(np.size(image_files))):
                    paths = []
                    portion = os.path.splitext(image_files[j])
                    newcamname = portion[0] + '.txt'
                    newdepthname = portion[0] + gt_fext

                    # ref image
                    ref_image_path = os.path.join(os.path.join(image_folder, ('%d' % ref_ind)),
                                                  image_files[j]).replace("\\", "/")
                    ref_cam_path = os.path.join(os.path.join(cam_folder, ('%d' % ref_ind)), newcamname).replace(
                        "\\", "/")

                    paths.append(ref_image_path)
                    paths.append(ref_cam_path)

                    # view images
                    for view in range(view_cnts - 1):
                        # print(image_folder)
                        view_ind = view_inds[view]  # selected view image
                        view_image_path = os.path.join(os.path.join(image_folder, ('%d' % view_ind)),
                                                       image_files[j]).replace("\\", "/")

                        view_cam_path = os.path.join(os.path.join(cam_folder, ('%d' % view_ind)),
                                                     newcamname).replace("\\", "/")
                        paths.append(view_image_path)
                        paths.append(view_cam_path)

                    # depth path
                    depth_image_path = os.path.join(os.path.join(depth_folder, ('%d' % ref_ind)),
                                                    newdepthname).replace("\\", "/")
                    paths.append(depth_image_path)
                    sample_list.append((sat_name, view_cnts, paths))

        return sample_list

    def dtu_list(self, data_folder, set_name, gt_fext='.pfm'):
        sample_list = []

        # image index
        train_cluster_path = data_folder + '/index.txt'
        data_cluster = open(train_cluster_path).read().split()

        # pair
        view_pair_path = data_folder + '/pair.txt'
        metas = []

        with open(view_pair_path) as f:
            cluster_num = int(f.readline().rstrip())
            for idx in range(cluster_num):
                ref_index = int(f.readline().rstrip())
                src_indexs = [int(x) for x in f.readline().rstrip().split()[1::2]]
                for light_idx in range(7):
                    metas.append((light_idx, ref_index, src_indexs))

        # for each data scene
        for i in data_cluster:
            image_folder = os.path.join(data_folder, ('Rectified/%s_train' % i)).replace("\\", "/")
            cam_folder = os.path.join(data_folder, 'Cameras').replace("\\", "/")
            depth_folder = os.path.join(data_folder, ('Depths/%s' % i)).replace("\\", "/")

            for idx in range(len(metas)):
                paths = []
                light_idx, ref_index, src_indexs = metas[idx]
                view_cnts = min(self.view_num, len(src_indexs) + 1)
                src_view_ids = src_indexs[:view_cnts - 1]

                ref_image_path = os.path.join(image_folder,
                                              'rect_{:0>3}_{}_r5000.png'.format(ref_index + 1, light_idx)).replace("\\",
                                                                                "/")
                ref_cam_path = os.path.join(cam_folder, '{:0>8}_cam.txt'.format(ref_index)).replace("\\", "/")
                depth_image_path = os.path.join(depth_folder, 'depth_map_{:0>4}.pfm'.format(ref_index)).replace("\\",
                                                                                "/")

                paths.append(ref_image_path)
                paths.append(ref_cam_path)

                for i, vid in enumerate(src_view_ids):
                    view_image_path = os.path.join(image_folder,
                                                   'rect_{:0>3}_{}_r5000.png'.format(vid + 1, light_idx)).replace("\\",
                                                                                "/")
                    view_cam_path = os.path.join(cam_folder, '{:0>8}_cam.txt'.format(vid)).replace("\\", "/")
                    paths.append(view_image_path)
                    paths.append(view_cam_path)

                paths.append(depth_image_path)
                sample_list.append((set_name, view_cnts, paths))

        return sample_list

    def BlendedMVS_list(self, data_folder, set_name, gt_fext='.pfm'):
        sample_list = []

        # image index
        train_cluster_path = data_folder + '/index.txt'
        data_cluster = open(train_cluster_path).read().split()

        # for each data scene
        for i in data_cluster:
            image_folder = os.path.join(data_folder, ('%s/blended_images' % i)).replace("\\", "/")
            cam_folder = os.path.join(data_folder, '%s/cams' % i).replace("\\", "/")
            depth_folder = os.path.join(data_folder, ('%s/rendered_depth_maps' % i)).replace("\\", "/")
            view_pair_path = os.path.join(data_folder, '%s/cams/pair.txt' % i).replace("\\", "/")

            metas = []
            with open(view_pair_path) as f:
                cluster_num = int(f.readline().rstrip())
                for idx in range(cluster_num):
                    ref_index = int(f.readline().rstrip())
                    src_indexs = [int(x) for x in f.readline().rstrip().split()[1::2]]
                    metas.append((ref_index, src_indexs))

            for idx in range(len(metas)):
                paths = []
                ref_index, src_indexs = metas[idx]
                view_cnts = min(self.view_num, len(src_indexs) + 1)
                src_view_ids = src_indexs[:view_cnts - 1]

                ref_image_path = os.path.join(image_folder, '{:0>8}.jpg'.format(ref_index)).replace("\\", "/")
                ref_cam_path = os.path.join(cam_folder, '{:0>8}_cam.txt'.format(ref_index)).replace("\\", "/")
                depth_image_path = os.path.join(depth_folder, '{:0>8}.pfm'.format(ref_index)).replace("\\", "/")
                paths.append(ref_image_path)
                paths.append(ref_cam_path)

                for i, vid in enumerate(src_view_ids):
                    view_image_path = os.path.join(image_folder, '{:0>8}.jpg'.format(vid)).replace("\\", "/")
                    view_cam_path = os.path.join(cam_folder, '{:0>8}_cam.txt'.format(vid)).replace("\\", "/")
                    paths.append(view_image_path)
                    paths.append(view_cam_path)

                paths.append(depth_image_path)
                sample_list.append((set_name, view_cnts, paths))

        return sample_list

    def ObliqueWhu_list(self, data_folder, set_name, gt_fext='.exr'):
        sample_list = []

        # image index
        train_cluster_path = data_folder + '/index.txt'
        data_cluster = open(train_cluster_path).read().split()

        def read_name_list(path):
            names_list = {}
            cluster_list = open(path).read().split()
            total_num = int(cluster_list[0])

            for i in range(total_num):
                index = int(cluster_list[i * 3 + 1])  # index
                name = cluster_list[i * 3 + 2]
                names_list[index] = name
            return names_list

        # for each data scene
        for i in data_cluster:
            image_folder = os.path.join(data_folder, ('%s/images' % i)).replace("\\", "/")
            cam_folder = os.path.join(data_folder, '%s/cams' % i).replace("\\", "/")
            depth_folder = os.path.join(data_folder, ('%s/depths' % i)).replace("\\", "/")
            view_pair_path = os.path.join(data_folder, '%s/info/viewpair.txt' % i).replace("\\", "/")
            map_index_path = os.path.join(data_folder, '%s/info/image_path.txt' % i).replace("\\", "/")
            map_index_dict = read_name_list(map_index_path)

            metas = []
            with open(view_pair_path) as f:
                cluster_num = int(f.readline().rstrip())
                for idx in range(cluster_num):
                    ref_index = int(f.readline().rstrip())
                    src_indexs = [int(x) for x in f.readline().rstrip().split()[1::2]]
                    metas.append((ref_index, src_indexs))

            for idx in range(len(metas)):
                paths = []
                ref_index, src_indexs = metas[idx]
                view_cnts = min(self.view_num, len(src_indexs) + 1)
                src_view_ids = src_indexs[:view_cnts - 1]

                ref_name = map_index_dict[ref_index]
                ref_image_path = os.path.join(image_folder, '{}.png'.format(ref_name)).replace("\\", "/")
                ref_cam_path = os.path.join(cam_folder, '{}.txt'.format(ref_name)).replace("\\", "/")
                depth_image_path = os.path.join(depth_folder, ref_name+gt_fext).replace("\\", "/")
                paths.append(ref_image_path)
                paths.append(ref_cam_path)

                for i, vid in enumerate(src_view_ids):
                    src_name = map_index_dict[vid]
                    view_image_path = os.path.join(image_folder, '{}.png'.format(src_name)).replace("\\", "/")
                    view_cam_path = os.path.join(cam_folder, '{}.txt'.format(src_name)).replace("\\", "/")
                    paths.append(view_image_path)
                    paths.append(view_cam_path)

                paths.append(depth_image_path)
                sample_list.append((set_name, view_cnts, paths))

        return sample_list

    def __len__(self):
        return len(self.sample_list)

    def tr_read_whu_cam(self, file, interval_scale=1):
        # read camera txt file
        cam = np.zeros((2, 4, 4), dtype=np.float32)
        extrinsics = np.zeros((4, 4), dtype=np.float32)
        pera = np.zeros((1, 13), dtype=np.float32)
        words = open(file).read().split()
        # read extrinsic
        for i in range(0, 4):
            for j in range(0, 4):
                extrinsic_index = 4 * i + j + 1
                extrinsics[i][j] = words[extrinsic_index]  # Twc

        # if cam ori is XrightYup
        O = np.eye(3, dtype=np.float32)
        O[1, 1] = -1
        O[2, 2] = -1
        R = extrinsics[0:3, 0:3]
        R2 = np.matmul(R, O)
        extrinsics[0:3, 0:3] = R2

        extrinsics = np.linalg.inv(extrinsics)  # Tcw
        cam[0, :, :] = extrinsics

        for i in range(0, 13):
            pera[0][i] = words[17 + i]

        f = pera[0][0]
        x0 = pera[0][1] # whu
        y0 = pera[0][2]

        # K Photogrammetry system XrightYup
        cam[1][0][0] = f
        cam[1][1][1] = f
        cam[1][0][2] = x0
        cam[1][1][2] = y0
        cam[1][2][2] = 1

        # depth range
        cam[1][3][0] = np.float32(pera[0][3])  # start
        cam[1][3][1] = np.float32(pera[0][5] * interval_scale)  # interval
        cam[1][3][3] = np.float32(pera[0][4])  # end

        acturald = int((cam[1][3][3] - cam[1][3][0]) / cam[1][3][1] / 32 + 1) * 32

        cam[1][3][2] = acturald

        location = words[23:30]
        return cam, location

    def tr_read_dtu_cam(self, file, interval_scale=1):
        # read camera txt file
        cam = np.zeros((2, 4, 4), dtype=np.float32)
        intrinsics = np.zeros((3, 3), dtype=np.float32)
        extrinsics = np.zeros((4, 4), dtype=np.float32)
        pera = np.zeros((1, 13), dtype=np.float32)
        words = open(file).read().split()
        # read extrinsic
        for i in range(0, 4):
            for j in range(0, 4):
                extrinsic_index = 4 * i + j + 1
                extrinsics[i][j] = words[extrinsic_index]  # Tcw

        cam[0, :, :] = extrinsics

        for i in range(0, 3):
            for j in range(0, 3):
                intrinsic_index = 3 * i + j + 18
                intrinsics[i][j] = words[intrinsic_index]

        cam[1, 0:3, 0:3] = intrinsics

        # depth range
        cam[1][3][0] = np.float32(words[27])  # start
        cam[1][3][1] = np.float32(words[28] * interval_scale)  # interval
        cam[1][3][3] = np.float32(cam[1][3][0] + cam[1][3][1] * 192)  # end

        acturald = int((cam[1][3][3] - cam[1][3][0]) / cam[1][3][1] / 32 + 1) * 32

        cam[1][3][2] = acturald

        location = [0, 0, 0, 0]
        return cam, location

    def tr_read_blendedmvs_cam(self, file, interval_scale=1):
        # read camera txt file
        cam = np.zeros((2, 4, 4), dtype=np.float32)
        intrinsics = np.zeros((3, 3), dtype=np.float32)
        extrinsics = np.zeros((4, 4), dtype=np.float32)
        pera = np.zeros((1, 13), dtype=np.float32)
        words = open(file).read().split()
        # read extrinsic
        for i in range(0, 4):
            for j in range(0, 4):
                extrinsic_index = 4 * i + j + 1
                extrinsics[i][j] = words[extrinsic_index]  # Tcw

        cam[0, :, :] = extrinsics

        for i in range(0, 3):
            for j in range(0, 3):
                intrinsic_index = 3 * i + j + 18
                intrinsics[i][j] = words[intrinsic_index]

        cam[1, 0:3, 0:3] = intrinsics

        # depth range
        cam[1][3][0] = np.float32(words[27])  # start
        cam[1][3][1] = np.float32(words[28]) * np.float32(interval_scale)  # interval
        cam[1][3][2] = np.float32(words[29]) / np.float32(interval_scale)  # sample number
        cam[1][3][3] = np.float32(words[30])  # end

        location = [0, 0, 0, 0]
        return cam, location

    def tr_read_obliquewhu_cam(self, file, interval_scale=1):
        # read camera txt file
        cam = np.zeros((2, 4, 4), dtype=np.float32)
        intrinsics = np.zeros((3, 3), dtype=np.float32)
        extrinsics = np.zeros((4, 4), dtype=np.float32)
        words = open(file).read().split()
        # read extrinsic
        for i in range(0, 4):
            for j in range(0, 4):
                extrinsic_index = 4 * i + j + 2
                extrinsics[i][j] = words[extrinsic_index]  # Twc

        # if cam ori is XrightYup
        O = np.eye(3, dtype=np.float32)
        O[1, 1] = -1
        O[2, 2] = -1
        R = extrinsics[0:3, 0:3]
        R2 = np.matmul(R, O)
        extrinsics[0:3, 0:3] = R2

        extrinsics = np.linalg.inv(extrinsics)  # Tcw
        cam[0, :, :] = extrinsics

        for i in range(0, 3):
            for j in range(0, 3):
                intrinsic_index = 3 * i + j + 18
                intrinsics[i][j] = words[intrinsic_index]
        cam[1, 0:3, 0:3] = intrinsics

        # depth range
        cam[1][3][0] = np.float32(words[27]) # start
        cam[1][3][3] = np.float32(words[28]) # end
        cam[1][3][1] = np.float32(words[29]) * interval_scale  # interval
        acturald = int((cam[1][3][3] - cam[1][3][0]) / cam[1][3][1] / 32 + 1) * 32
        cam[1][3][2] = acturald

        location = [0, 0, 0, 0]

        return cam, location

    def read_img(self, filename):
        img = Image.open(filename)
        return img

    def read_depth(self, filename, set_name):
        if set_name =='whu_mvs':
            depimg = imread(filename) # read png depth file
            depth_image = (np.float32(depimg) / 64.0)  # WHU-MVS dataset

        elif set_name == 'dtu' or set_name == 'BlendedMVS':
            depth_image = np.array(read_pfm(filename)[0], dtype=np.float32)  # read pfm depth file

        elif set_name == 'whu_omvs':
            depth_image = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
            mask_path = filename.replace("depths", "masks")
            mask_path = mask_path.replace(".exr", ".png")
            mask_image = np.array(cv2.imread(mask_path, cv2.COLOR_BGR2GRAY)) / 255.
            mask_image = mask_image < 0.5
            depth_image[mask_image] = 0

        else:
            raise Exception("{}? Not implemented yet!".format(set_name))

        return np.array(depth_image)

    def center_image(self, img, mode='mean'):
        # normalize image input
        if mode == 'standard':
            np_img = np.array(img, dtype=np.float32) / 255.

        elif mode == 'mean':
            img_array = np.array(img)
            img = img_array.astype(np.float32)
            var = np.var(img, axis=(0, 1), keepdims=True)
            mean = np.mean(img, axis=(0, 1), keepdims=True)
            np_img = (img - mean) / (np.sqrt(var) + 0.00000001)

        else:
            raise Exception("{}? Not implemented yet!".format(mode))

        return np_img

    def __getitem__(self, idx):

        sat_name, view_cnts, paths = self.sample_list[idx]

        # read input data
        outimage = None
        outcam = None
        outlocation = None
        centered_images = []
        proj_matrices = []

        # depth
        depth_image = self.read_depth(os.path.join(paths[2 * view_cnts]), sat_name)

        for view in range(view_cnts):
            # images
            if self.mode == "train":
                image = image_augment(self.read_img(paths[2 * view]))
            else:
                image = self.read_img(paths[2 * view])
            image = np.array(image)

            # cameras
            if sat_name == 'dtu':
                cam, _ = self.tr_read_dtu_cam(paths[2 * view + 1], self.interval_scale)
            elif sat_name == 'BlendedMVS':
                cam, _ = self.tr_read_blendedmvs_cam(paths[2 * view + 1], self.interval_scale)
            elif sat_name == 'whu_omvs':
                cam, _ = self.tr_read_obliquewhu_cam(paths[2 * view + 1], self.interval_scale)
            elif sat_name == 'whu_mvs':
                cam, _ = self.tr_read_whu_cam(paths[2 * view + 1], self.interval_scale)
            else:
                raise Exception("{}? Not implemented yet!".format(sat_name))

            if view == 0:
                outimage = image
                outcam = cam
                depth_min = outcam[1][3][0]
                depth_interval = outcam[1][3][1]
                depth_max = outcam[1][3][3]

            # multiply intrinsics and extrinsics to get projection matrix
            extrinsics = cam[0, :, :]
            intrinsics = cam[1, 0:3, 0:3]
            proj_mat = extrinsics.copy()
            proj_mat[:3, :4] = np.matmul(intrinsics, proj_mat[:3, :4])

            proj_matrices.append(proj_mat)
            centered_images.append(self.center_image(image))

        centered_images = np.stack(centered_images).transpose([0, 3, 1, 2])
        proj_matrices = np.stack(proj_matrices)

        # Depth
        depth_values = np.array([depth_min, depth_max, depth_interval], dtype=np.float32)
        mask = np.float32((depth_image >= depth_min) * 1.0) * np.float32((depth_image <= depth_max) * 1.0)

        h, w = depth_image.shape
        depth_ms = {
            "stage1": cv2.resize(depth_image, (w//4, h//4), interpolation=cv2.INTER_NEAREST),
            "stage2": cv2.resize(depth_image, (w//2, h//2), interpolation=cv2.INTER_NEAREST),
            "stage3": depth_image,
        }
        mask_ms = {
            "stage1": cv2.resize(mask, (w // 4, h // 4), interpolation=cv2.INTER_NEAREST),
            "stage2": cv2.resize(mask, (w // 2, h // 2), interpolation=cv2.INTER_NEAREST),
            "stage3": mask,
        }
        # ms proj_mats
        stage2_pjmats = proj_matrices.copy()
        stage2_pjmats[:, :2, :] = proj_matrices[:, :2, :] / 2
        stage3_pjmats = proj_matrices.copy()
        stage3_pjmats[:, :2, :] = proj_matrices[:, :2, :] / 4

        proj_matrices_ms = {
            "stage1": stage3_pjmats,
            "stage2": stage2_pjmats,
            "stage3": proj_matrices
        }

        if sat_name == 'whu_omvs':
            name = os.path.splitext(os.path.basename(paths[0]))[0]
            vid = os.path.dirname(paths[0]).split("/")[-1]
        else:
            name = os.path.splitext(os.path.basename(paths[0]))[0]
            vid = os.path.dirname(paths[0]).split("/")[-2]

        return {"imgs": centered_images,
                "proj_matrices": proj_matrices_ms,
                "depth": depth_ms,
                "mask": mask_ms,
                "depth_values": depth_values,
                "depth_interval": depth_interval,
                "outimage": outimage,
                "outcam": outcam,
                "out_name": name,
                "out_view": vid}

if __name__ == "__main__":

    pass