# https://github.com/haosulab/ManiSkill2022/blob/plb/mpm/simulator.py
import numpy as np
import torch

from mpm.types import (
    ivec3,
    array,
    float32,
    vec3,
    quat,
    mat3,
    lib,
)

grid_lower_default = np.array([[1e6, 1e6, 1e6]], dtype=np.int32)
grid_lower_zero = np.array([[0., 0., 0.]], dtype=np.int32)

def rigid_body_motion(states, actions):
    import torch
    # state: (B, 7)
    # action: (T, B, 6)
    T, B = actions.shape[:2]
    pos, q = states

    pos = pos[None, :].expand(T, -1, -1).reshape(-1, 3)
    q = q[None, :].expand(T, -1, -1).reshape(-1, 4)
    actions = actions.reshape(-1, 6)
    rot = actions[:, 3:]
    # better if we can move this part to C++ code
    w = torch.sqrt((rot * rot).sum(axis=-1, keepdims=True) + 1e-16)
    quat = torch.cat((torch.cos(w / 2), (rot / torch.clamp(w, 1e-7, 1e9)) * torch.sin(w / 2)), 1)

    next_pos = pos + actions[:, :3]
    # terms = q.outer_product(quat) # we should use w x q
    terms = q[:, :, None] * quat[:, None, :]
    w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]
    x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]
    y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]
    z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]
    out = torch.stack([w, x, y, z], 1)
    next_rot = out / torch.linalg.norm(out, dim=-1, keepdims=True)
    return next_pos.reshape(T, B, 3), next_rot.reshape(T, B, 4)

class TempState:
    # store temporary grid information
    def __init__(self, n_particles, grid_dim, n_bodies):
        grid_length = grid_dim[0] * grid_dim[1] * grid_dim[2]
        self.grid_m = array(dtype=float, length=grid_length)
        self.grid_v_in = array(dtype=vec3, length=grid_length)  # v_in
        self.grid_v_out = array(dtype=vec3, length=grid_length)  # v_out

        self.grid_m_grad = array(dtype=float, length=grid_length)
        self.grid_v_in_grad = array(dtype=vec3, length=grid_length)
        self.grid_v_out_grad = array(dtype=vec3, length=grid_length)

        self.grid_body_v_in = array(dtype=vec3, length=grid_length * (n_bodies + 1))

        self.F = array(dtype=mat3, length=n_particles)  # deformation gradient
        self.U = array(dtype=mat3, length=n_particles)  # deformation gradient
        self.V = array(dtype=mat3, length=n_particles)  # deformation gradient
        self.sig = array(dtype=vec3, length=n_particles)  # deformation gradient

        self.F_grad = array(dtype=mat3, length=n_particles)  # deformation gradient
        self.U_grad = array(dtype=mat3, length=n_particles)  # deformation gradient
        self.V_grad = array(dtype=mat3, length=n_particles)  # deformation gradient
        self.sig_grad = array(dtype=vec3, length=n_particles)  # deformation gradient

        self.grid_lower = array(dtype=ivec3, length=1)

    def clear(self, stream=None):
        self.grid_m.zero(stream)
        self.grid_v_in.zero(stream)
        self.grid_v_out.zero(stream)

    def clear_grad(self, stream=None):
        self.grid_v_in_grad.zero(stream)
        self.grid_v_out_grad.zero(stream)
        self.grid_m_grad.zero(stream)

        self.sig_grad.zero(stream)
        self.F_grad.zero(stream)
        self.U_grad.zero(stream)
        self.V_grad.zero(stream)

class State:
    def __init__(self, n_particles, n_bodies):
        self.n_particles = int(n_particles)
        self.n_bodies = n_bodies

        self.x = array(dtype=vec3, length=n_particles)
        self.v = array(dtype=vec3, length=n_particles)
        self.F = array(dtype=mat3, length=n_particles)
        self.C = array(dtype=mat3, length=n_particles)

        self.x_grad = array(dtype=vec3, length=n_particles)
        self.v_grad = array(dtype=vec3, length=n_particles)
        self.F_grad = array(dtype=mat3, length=n_particles)
        self.C_grad = array(dtype=mat3, length=n_particles)

        if n_bodies:
            self.body_pos = array(dtype=vec3, length=n_bodies)
            self.body_rot = array(dtype=quat, length=n_bodies)

            self.body_pos_grad = array(dtype=vec3, length=n_bodies)
            self.body_rot_grad = array(dtype=quat, length=n_bodies)

    def reset(self):
        eye = np.zeros((self.model.n_particles, 3, 3), dtype=np.float32)
        eye[:] = np.eye(3)
        self.x.zero()
        self.v.zero()
        self.C.zero()
        self.F.upload(eye.reshape(self.F.shape))

    def get_state(self, n=None):
        x = self.x.download(n)
        v = self.v.download(n)
        F = self.F.download(n)
        C = self.C.download(n)
        return x, v, F.reshape(-1, 3, 3), C.reshape(-1, 3, 3), \
            self.body_pos.download(), self.body_rot.download()

    def set_state(self, state):
        self.x.upload(state[0])
        self.v.upload(state[1])
        self.F.upload(state[2].reshape(-1, 9))
        self.C.upload(state[3].reshape(-1, 9))

        pos, rot = state[4:]
        self.body_pos.upload(pos)
        self.body_rot.upload(rot)

    def clear_grad(self, stream):
        # TODO: should be asynchronized later after back-propogation..
        self.x_grad.zero_async(stream)
        self.v_grad.zero_async(stream)
        self.F_grad.zero_async(stream)
        self.C_grad.zero_async(stream)

        if self.n_bodies:
            self.body_pos_grad.zero_async(stream)
            self.body_rot_grad.zero_async(stream)

from tools import Configurable

class MPMSimulator(Configurable):
    def __init__(
            self,
            n_bodies,
            cfg=None,
            ground_friction=0.,  # Currently I only support 0 now.
            gravity=(0., -1., 0.),  # 9.8
            n_particles=20000,
            dx=1. / 64,
            dt=0.0001,
            grid_size=(1., 1., 1.),
            max_steps=30,
            substeps=20,

            yield_stress=30.,
            vol=(1. / 64 / 2) ** 2,  # based on plasticinelab's setup
            mass=(1. / 64 / 2) ** 2,
            E=5000.,
            nu=0.2,
    ):
        super().__init__()
        # self.cfg = cfg
        # self.process_actors(actors)
        self.dx = dx
        self.dt = dt
        self.max_particles = n_particles
        self.n_particles = n_particles
        self.n_bodies = n_bodies
        self.substeps = substeps

        grid_dim = np.ceil(np.array(grid_size) / self.dx / 4).astype(int) * 4
        self.grid_length = grid_dim[0] * grid_dim[1] * grid_dim[2]

        self.dx = float(dx)
        self.inv_dx = 1.0 / self.dx
        self.dt = float(dt)
        self.grid_dim = ivec3(grid_dim[0], grid_dim[1], grid_dim[2])

        self.gravity = array(dtype=vec3, length=1)
        self.particle_vol = array(dtype=float32, length=self.n_particles)
        self.particle_mass = array(dtype=float32, length=self.n_particles)
        self.particle_mu_lam_yield = array(dtype=vec3, length=self.n_particles)
        self.particle_color = array(dtype=int, length=self.n_particles)

        self.particle_ids = array(dtype=int, length=self.n_particles)

        self.dists = array(dtype=float, length=self.n_particles * n_bodies)
        self.dists_grad = array(dtype=float, length=self.n_particles * n_bodies)

        if n_bodies:
            self.body_type_mu_softness_round = array(dtype=quat, length=n_bodies)
            self.body_args = array(dtype=quat, length=n_bodies)  # four args now ..

        self.temp = TempState(self.n_particles, grid_dim, self.n_bodies)
        self.max_steps = max_steps
        self.states = [State(self.n_particles, n_bodies) for _ in range(max_steps)]

        self.stream0 = lib.cuda_stream_create()  # compute stream
        self.stream1 = lib.cuda_stream_create()  # data stream

        mu = E / (2 * (1 + nu))
        lam = E * nu / ((1 + nu) * (1 - 2 * nu))
        self.init_particles(
            np.zeros(self.n_particles) + vol,
            np.zeros(self.n_particles) + mass,
            np.zeros((self.n_particles, 3)) + np.array([mu, lam, yield_stress]),
        )
        self.n_bodies = n_bodies

        self.object_id = None

        self.torch_scale = None

        self.n_grid = grid_dim[0]

        lib.print_memory_info()

    def clear_grad(self, max_steps):
        for i in range(max_steps):
            self.states[i].clear_grad()

    def get_state(self, index):
        # to be consistent with the PlasticineLab design
        state = self.states[index].get_state(self.n_particles)
        return tuple(state[:4]) + tuple([np.r_[a, b] for a, b in zip(state[4], state[5])])

    def set_state(self, index, state):
        state = [np.float32(i) for i in state]
        self.n_particles = int(len(state[0]))

        pos = [i[:3] for i in state[4:]]
        rot = [i[3:] for i in state[4:]]
        pos = np.float32(pos)
        rot = np.float32(rot)
        state = state[:4] + [pos, rot]
        # for i in range(4):
        #   empty = np.zeros((self.max_particles, *state[i].shape[1:]))
        #   empty[:self.n_particles] = state[i]
        #   state[i] = empty
        self.states[index].set_state(state)
        self.cur = index

    def set_object_id(self, object_id):
        self.object_id = object_id
        assert len(object_id) == self.n_particles, f"{len(object_id)} {self.n_particles}"
        self.particle_ids.upload(self.object_id)

    def get_object_id(self, device):
        assert device == 'numpy'
        # print(self.object_id.shape, self.n_particles)
        return self.object_id

    def set_color(self, inp):
        color = self.particle_color.download(self.n_particles)
        color[:] = inp
        self.particle_color.upload(color)

    def set_softness(self, softness):
        pp = self.body_type_mu_softness_round.download()
        pp[:, 2] = softness
        self.body_type_mu_softness_round.upload(pp)

    def get_softness(self):
        pp = self.body_type_mu_softness_round.download()
        return pp[:, 2]

    def _initialize_buffer(self, device='numpy', dims=None, dtype=None):
        dims = dims or (3,)
        if device == 'numpy':
            dtype = dtype or np.float32
            x = np.zeros((self.n_particles, *dims), dtype=dtype)
        else:
            import torch
            dtype = dtype or torch.float32
            x = torch.zeros((self.n_particles, *dims), dtype=torch.float32, device=device)
        return x

    def get_x(self, index, device='numpy'):
        return self.states[index].x.download(self.n_particles, device=device)

    def get_v(self, index, device='numpy'):
        return self.states[index].v.download(self.n_particles, device=device)

    def get_dists(self, f, grad=None, device='cuda:0'):
        state = self.states[f]
        if grad is None:
            need_grad = False
        else:
            need_grad = True
            assert grad.shape == (self.n_particles, self.n_bodies)
            grad = grad.reshape(-1).detach().cpu().numpy()
            self.dists_grad.upload(grad)

        lib.compute_dist(
            state.x.data_ptr,
            state.body_pos.data_ptr,
            state.body_rot.data_ptr,
            self.body_type_mu_softness_round.data_ptr,
            self.body_args.data_ptr,
            self.dists.data_ptr,
            self.n_bodies,
            state.x_grad.data_ptr,
            state.body_pos_grad.data_ptr,
            state.body_rot_grad.data_ptr,
            self.dists_grad.data_ptr,
            int(need_grad),
            self.n_particles,
            self.stream0
        )
        if grad is None:
            return self.dists.download(self.n_particles * self.n_bodies, device=device).reshape(-1, self.n_bodies)

    def compute_grid_mass(self, f, id=-1, device='numpy', backward_grad=None):
        need_grad = backward_grad is not None
        # assert not need_grad

        if isinstance(f, int):
            n_particles = self.n_particles
            state = self.states[f]
        else:
            if isinstance(f, torch.Tensor):
                f = f.detach().cpu().numpy()
            assert f.shape[-1] == 3 and id == -1
            state = self.states[-1]
            state.x.upload(np.float32(f))
            n_particles = len(f)
            if need_grad:
                state.clear_grad(self.stream0)
            # assert not need_grad

        if backward_grad is not None:
            self.temp.grid_m_grad.upload(backward_grad.reshape(-1).detach().cpu().numpy(), strict=True)

        self.temp.grid_m.zero()
        # raise NotImplementedError
        lib.particle2mass(state.x.data_ptr, self.particle_mass.data_ptr, self.temp.grid_lower.data_ptr, self.grid_dim,
                          self.dx, self.inv_dx, self.temp.grid_m.data_ptr, self.temp.grid_m_grad.data_ptr,
                          state.x_grad.data_ptr, self.particle_ids.data_ptr, id, need_grad, n_particles, self.stream0)

        if backward_grad is None:
            return self.temp.grid_m.download(device=device).reshape(*self.grid_dim)

    def compute_svd(self, index=0, device=None, backward_grad=None):
        raise NotImplementedError

    def get_tool_state(self, index=0, device=None):
        pos = self.states[index].body_pos.download()
        rot = self.states[index].body_rot.download()
        state = [i for i in np.concatenate((pos, rot), axis=1)]
        if device == 'numpy':
            return state
        else:
            return [torch.tensor(i).to(device) for i in state]

    def set_tool_state(self, index, pose):
        pose = np.array(pose)
        assert pose.shape == (self.n_bodies, 7), f"{pose.shape}, {self.n_bodies}"
        pos = pose[:, :3]
        rot = pose[:, 3:7]
        self.states[index].body_pos.upload(pos)
        self.states[index].body_rot.upload(rot)

    def get_action_scales(self):
        return np.array(self.action_scales)

    def init_particles(self, vol, mass, mu_lam_yield):
        assert vol.shape == (self.n_particles,)
        assert mass.shape == (self.n_particles,)
        assert mu_lam_yield.shape == (self.n_particles, 3)

        self.particle_vol.upload(vol)
        self.particle_mass.upload(mass)
        self.particle_mu_lam_yield.upload(mu_lam_yield)

        self.gravity.upload(np.float32([np.array(self._cfg.gravity) * 30]))
        self.temp.grid_lower.upload(grid_lower_zero)

    def init_bodies(self, types, softness, mu, round, args, action_scales, pos=None, rot=None):
        assert len(mu) == len(args) == self.n_bodies
        mu = np.float32(mu)
        softness = np.float32(softness)
        types = np.float32(types)
        round = np.float32(round)
        args = np.float32(args)

        self.body_type_mu_softness_round.upload(
            np.stack((types, mu, softness, round), 1).reshape(-1, 4))
        self.body_args.upload(args)
        self.action_scales = action_scales

        if pos is not None:
            self.states[0].body_pos.upload(np.array(pos).reshape(self.n_bodies, 3))
        if rot is not None:
            self.states[0].body_rot.upload(np.array(rot).reshape(self.n_bodies, 4))

    def sync(self):
        lib.cuda_stream_sync(self.stream1)
        lib.cuda_stream_sync(self.stream0)

    def __del__(self):
        if lib is not None:
            lib.cuda_stream_destroy(self.stream0)
            lib.cuda_stream_destroy(self.stream1)

    def nan_check(self, i, message=""):
        if np.isnan(self.states[i].particle_x.download()).any():
            print("x nan", message)
        if np.isnan(self.states[i].particle_v.download()).any():
            print("v nan", message)
        if np.isnan(self.states[i].particle_F.download()).any():
            print("F nan", message)
        if np.isnan(self.states[i].particle_C.download()).any():
            print("C nan", message)

    def compute_grid_lower(self, state: State, temp: TempState):
        lib.compute_grid_lower(
            state.x.data_ptr,
            self.dx,
            self.inv_dx,
            temp.grid_lower.data_ptr,
            self.n_particles,
            self.stream0,
        )

    def compute_svd(self, state: State, temp: TempState):
        lib.compute_svd(
            state.F.data_ptr,
            state.C.data_ptr,
            temp.F.data_ptr,
            temp.U.data_ptr,
            temp.V.data_ptr,
            temp.sig.data_ptr,
            self.dt,
            self.n_particles,
            self.stream0,
        )

    def compute_svd_grad(self, state: State, temp: TempState):
        lib.compute_svd_grad(
            state.F.data_ptr,
            state.C.data_ptr,
            temp.U.data_ptr,
            temp.V.data_ptr,
            temp.sig.data_ptr,

            temp.F_grad.data_ptr,
            temp.U_grad.data_ptr,
            temp.V_grad.data_ptr,
            temp.sig_grad.data_ptr,
            state.F_grad.data_ptr,
            state.C_grad.data_ptr,

            self.dt,
            self.n_particles,
            self.stream0,
        )

    def p2g(self, state1: State, temp: TempState, state2: State):
        lib.p2g(
            state1.x.data_ptr, state1.v.data_ptr, self.particle_mass.data_ptr, self.particle_vol.data_ptr,
            temp.F.data_ptr, temp.U.data_ptr, temp.sig.data_ptr, temp.V.data_ptr, state1.C.data_ptr,
            self.particle_mu_lam_yield.data_ptr, temp.grid_lower.data_ptr, self.grid_dim, self.dx, self.inv_dx, self.dt,
            state2.F.data_ptr, temp.grid_v_in.data_ptr, temp.grid_m.data_ptr, self.n_particles, self.stream0
        )

    def p2g_grad(self, state1: State, temp: TempState, state2: State):
        lib.p2g_grad(
            state1.x.data_ptr, state1.v.data_ptr, self.particle_mass.data_ptr, self.particle_vol.data_ptr,
            temp.F.data_ptr, temp.U.data_ptr, temp.sig.data_ptr, temp.V.data_ptr, state1.C.data_ptr,
            self.particle_mu_lam_yield.data_ptr, temp.grid_lower.data_ptr, self.grid_dim, self.dx, self.inv_dx, self.dt,
            state2.F.data_ptr, temp.grid_v_in.data_ptr, temp.grid_m.data_ptr,

            state1.x_grad.data_ptr, state1.v_grad.data_ptr, temp.F_grad.data_ptr, state1.C_grad.data_ptr,
            temp.U_grad.data_ptr, temp.sig_grad.data_ptr, temp.V_grad.data_ptr, state2.F_grad.data_ptr,
            temp.grid_v_in_grad.data_ptr, temp.grid_m_grad.data_ptr,

            self.n_particles, self.stream0
        )

    def get_ground_friction(self):
        return float(self._cfg.ground_friction)

    def get_ground_height(self):
        return float(3.)

    def grid_op(self, state1: State, temp: TempState, state2: State):
        lib.grid_op_v2(
            temp.grid_m.data_ptr, temp.grid_v_in.data_ptr, temp.grid_body_v_in.data_ptr,
            temp.grid_lower.data_ptr, self.gravity.data_ptr,
            state1.body_pos.data_ptr, state1.body_rot.data_ptr,
            state2.body_pos.data_ptr, state2.body_rot.data_ptr,
            self.body_type_mu_softness_round.data_ptr, self.body_args.data_ptr,
            self.dx, self.inv_dx, self.dt,
            self.get_ground_friction(),
            temp.grid_v_out.data_ptr, self.grid_dim, self.n_bodies,
            self.stream0,
        )

    def grid_op_grad(self, state1: State, temp: TempState, state2: State):
        lib.grid_op_v2_grad(
            temp.grid_m.data_ptr, temp.grid_v_in.data_ptr, temp.grid_body_v_in.data_ptr,
            temp.grid_lower.data_ptr, self.gravity.data_ptr,
            state1.body_pos.data_ptr, state1.body_rot.data_ptr,
            state2.body_pos.data_ptr, state2.body_rot.data_ptr,
            self.body_type_mu_softness_round.data_ptr, self.body_args.data_ptr,

            temp.grid_m_grad.data_ptr, temp.grid_v_in_grad.data_ptr, state1.body_pos_grad.data_ptr,
            state1.body_rot_grad.data_ptr, state2.body_pos_grad.data_ptr, state2.body_rot_grad.data_ptr,

            self.dx, self.inv_dx, self.dt,
            self.get_ground_friction(),
            temp.grid_v_out.data_ptr,
            temp.grid_v_out_grad.data_ptr,
            self.grid_dim, self.n_bodies,
            self.stream0,
        )

    def g2p(self, state1: State, temp: TempState, state2: State):
        lib.g2p(
            state1.x.data_ptr, temp.grid_v_out.data_ptr, temp.grid_lower.data_ptr,
            self.dx, self.inv_dx, self.dt,
            self.grid_dim, state2.v.data_ptr, self.get_ground_height(),
            state2.C.data_ptr, state2.x.data_ptr, self.n_particles, self.stream0
        )
        # print(state2.x.download()[0])

    def g2p_grad(self, state1: State, temp: TempState, state2: State):
        lib.g2p_grad(
            state1.x.data_ptr, temp.grid_v_out.data_ptr, temp.grid_lower.data_ptr,
            self.dx, self.inv_dx, self.dt,
            self.grid_dim, state2.v.data_ptr, self.get_ground_height(),
            state2.C.data_ptr, state2.x.data_ptr, self.n_particles,

            state1.x_grad.data_ptr,
            temp.grid_v_out_grad.data_ptr,
            state2.v_grad.data_ptr,
            state2.C_grad.data_ptr,
            state2.x_grad.data_ptr,

            self.stream0
        )

    def set_pose(self, state, pos, rot, stream):
        if isinstance(pos, torch.Tensor):
            pos = pos.detach().cpu().numpy()
        if isinstance(rot, torch.Tensor):
            rot = rot.detach().cpu().numpy()
        state.body_pos.upload_async(pos, stream=stream)
        state.body_rot.upload_async(rot, stream=stream)

    def substep(self, f, clear_grad=False):
        temp, cur, next = self.temp, self.states[f], self.states[f + 1]
        self.temp.clear(self.stream0)
        # self.sync()
        self.compute_svd(cur, temp)
        self.p2g(cur, temp, next)
        self.grid_op(cur, temp, next)
        self.g2p(cur, temp, next)

        if clear_grad:
            next.clear_grad(self.stream1)  # async clear grad

    def substep_grad(self, f):
        temp, cur, next = self.temp, self.states[f], self.states[f + 1]
        self.temp.clear(self.stream0)
        self.temp.clear_grad(self.stream0)
        # self.sync()
        self.compute_svd(cur, temp)
        self.p2g(cur, temp, next)
        self.grid_op(cur, temp, next)

        self.g2p_grad(cur, temp, next)
        self.grid_op_grad(cur, temp, next)
        self.p2g_grad(cur, temp, next)
        self.compute_svd_grad(cur, temp)

    def download_pos_rot(self, cur, device):
        pos, rot = self.states[cur].body_pos.download(), self.states[cur].body_rot.download()
        pos, rot = torch.tensor(pos, device=device, dtype=torch.float32), torch.tensor(rot, device=device,
                                                                                dtype=torch.float32)
        return pos, rot

    # substeps 2000
    # pytorch batch
    # theta; pos_p, pos_q = [pos_p, rot_p] * [axis * angle, trans]

    def compute_forward_kinematics(self, f, action, pos_rot=None):
        cur = f
        device = 'cpu' if not isinstance(action, torch.Tensor) else action.device

        if pos_rot is None:
            pos, rot = self.download_pos_rot(cur, device)
        else:
            pos, rot = pos_rot

        if not isinstance(action, torch.Tensor):
            action = np.array(action, np.float32)
            action = torch.tensor(action, device=device, dtype=torch.float32)
        else:
            action = action.to(device)

        if self.torch_scale is None:
            self.torch_scale = torch.tensor(np.array(self.action_scales), device=device, dtype=torch.float32)
        else:
            self.torch_scale = self.torch_scale.to(device)
        action = action.reshape(-1, 6).clamp(-1., 1.) * self.torch_scale

        pos, rot = rigid_body_motion(
            (pos, rot),
            action[None, :].expand(self.substeps, -1, -1) * (
                    torch.arange(self.substeps, device=device)[:, None, None] + 1
            ) / self.substeps
        )
        return pos, rot, (pos, rot)

    def step(self, action, pos_rot=None):
        pos, rot, q_state = self.compute_forward_kinematics(self.cur, action, pos_rot=pos_rot)

        for f in range(0, self.substeps):
            self.set_pose(self.states[f + 1], pos[f], rot[f], self.stream0)
            self.substep(f)

        self.sync()
        self.states[0], self.states[self.substeps] = self.states[self.substeps], self.states[0]


import binascii as b

def hex2str(hex):
    return b.a2b_hex(hex).decode('utf-8')

'''
去掉字符串前后的空格、制表符
'''
def strStrip(s):
    return str(s).strip('\t').strip(' ')

'''
解析HBD请求得到的响应报文
'''
def parseHBD(value):
    value = strStrip(value)
    msgHeader = hex2str(value[0: 8])
    reportMask = value[8: 10]
    length = int(value[10: 12], 16)
    diviceType = value[12: 14]
    protocolVersion = str(value[14: 16])+"."+str(value[16: 18])
    firmwareVersion = str(value[18: 20]) + str(value[20: 22])

    bReportMask = bin(int(reportMask, 16)).replace('0b', '')
    if bReportMask[3] == 0:
        uniqueID = ''
        for i in range(8):
            index = 22+i*2
            uniqueID += str(int(value[index: index+2], 16))
    else:
        uniqueID = hex2str(value[22: 38])
    sendTime = str(int(value[38: 42], 16)) + str(int(value[42: 43], 16)) + str(int(value[43: 44], 16)) + str(int(value[44: 46], 16)) \
               + str(int(value[46: 47], 16)) + str(int(value[47: 48], 16)) + str(int(value[48: 50], 16)) + str(int(value[50: 52], 16))
    countNumber = value[52: 56]
    checksum = value[56: 60]
    tailCharacters = value[60: 64]
    result = "Message Header: "+msgHeader + "\nReport Mask: " + reportMask + "\nLength: " +\
    str(length) +"\nDevice Type: " + diviceType + "\nProtocol Version: " + protocolVersion + "\nFirmware Version: " +\
    firmwareVersion + "\nUnique ID: " + uniqueID + "\nSend Time: " + sendTime + "\nCount Number: " +\
    countNumber + "\nChecksum: " + checksum + "\nTail Characters: " + tailCharacters
    print(result)
    return result

'''
解析DAT请求得到的响应报文
'''
def parseDAT(value):
    value = strStrip(value)
    print(value)

'''
解析EVT请求得到的响应报文
'''
def parseEVT(value):
    value = strStrip(value)
    msgHeader = hex2str(value[0: 8])
    reportMask = value[8: 10]
    length = int(value[10: 12], 16)
    diviceType = value[12: 14]
    protocolVersion = str(value[14: 16]) + "." + str(value[16: 18])
    firmwareVersion = str(value[18: 20]) + str(value[20: 22])

'''
解析请求以及响应报文
'''
def parse(key, value):
    key = strStrip(key)
    key = key[1:4]
    resultFile = open('data/result.log','w')
    if 'HBD' == key:
        resultFile.write(parseHBD(value))
    elif 'EVT' == key:
        parseEVT(value)

    resultFile.flush()
    resultFile.close()

f = open("data/GV300N.log", 'r')
aesc = 0  # 区分报文日志的顺序，1代表正序，0代表倒序
hashMap = {}
value = ""

while True:
    line = f.readline()
    if not line:
        break
    # print(line)
    if aesc == 1:  # 顺序
        line = f.readline()
    else:  # 倒序
        if line.__contains__('HEX'):
            value = line[line.find('HEX:') + len("HEX:"):]
            # print(value)
        elif line.__contains__('ASC'):
            key = line[line.find('ASC:') + len("ASC:"):]
            # print(key)
            hashMap[key] = value

# print(count)
for key, value in hashMap.items():
    # print('key is %s,value is %s' % (key, value))
    parse(key, value)


"""
Combination of two Sky Segmentation Implementations

1. Sky Detection in Images for Solar Exposure Prediction : https://core.ac.uk/download/pdf/35462487.pdf
2. Sky Region Detection in a Single Image for Autonomous Ground Robot Navigation: https://journals.sagepub.com/doi/pdf/10.5772/56884

Algorithm inspired by Journal 2 and Channel extraction inspired by Journal 1

"""

import cv2 as cv
import numpy as np
from scipy import spatial
from matplotlib import pyplot as plt
import glob
import os

#global iteration value to track images
ITER = 0

# Optional Function in order to compare original images to their masked/sky_segmented versions
def CompareImgs(img_1,img_2):
    plt.subplots(1,2)
    plt.subplot(1,2,1)
    plt.imshow(img_1)
    plt.subplot(1,2,2)
    plt.imshow(img_2)
    plt.show()

# Iterates through image to create mask
def GenerateMask(b, image):
    mask = np.zeros((image.shape[0], image.shape[1], 1), dtype=np.uint8)
    for xx, yy in enumerate(b):
        mask[yy:, xx] = 255

    return mask

# Saves images based on order
def SaveImage(img):
    cv.imwrite(("output/image_{}.jpg".format(ITER)), img)

# Applies mask and saves image (Can optionally also uncomment compare_imgs in order to view images and their masks side-by-side)
def DisplayMask(b, image, color=[0, 0, 0]):
    result = cv.bitwise_and(image, image, mask=cv.bitwise_not(GenerateMask(b, image))) # Mask inverted in order to save images of segmented sky
    #compare_imgs(image, result)
    SaveImage(result)

# Image is first converted to grayscale and then
# calculates the gradient magnitude by combining two gradient images (horizontal and vertical)
# generated using sobel functions.
def ColorToGradient(image):
    gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
    return np.hypot(cv.Sobel(gray, cv.CV_64F, 1, 0), cv.Sobel(gray, cv.CV_64F, 0, 1))

# Energy Function extracted directly from Journal 2
def EnergyFunction(b_tmp, image):
    sky_mask = GenerateMask(b_tmp, image)

    ground = np.ma.array(image, mask=cv.cvtColor(cv.bitwise_not(sky_mask), cv.COLOR_GRAY2BGR)).compressed()
    sky = np.ma.array(image, mask=cv.cvtColor(sky_mask, cv.COLOR_GRAY2BGR)).compressed()

    ground.shape = (ground.size//3, 3)
    sky.shape = (sky.size//3, 3)

    sigma_g, mu_g = cv.calcCovarMatrix(ground, None, cv.COVAR_NORMAL | cv.COVAR_ROWS | cv.COVAR_SCALE)
    sigma_s, mu_s = cv.calcCovarMatrix(sky, None, cv.COVAR_NORMAL | cv.COVAR_ROWS | cv.COVAR_SCALE)

    y = 2

    return 1 / (
        (y * np.linalg.det(sigma_s) + np.linalg.det(sigma_g)) +
        (y * np.linalg.det(np.linalg.eig(sigma_s)[1]) +
            np.linalg.det(np.linalg.eig(sigma_g)[1]))
    )

# Detect position of horizon. Taken from Journal 2.
def CalcBorder(grad, t):
    sky_b = np.full(grad.shape[1], grad.shape[0])

    for x in range(grad.shape[1]):
        border_pos = np.argmax(grad[:, x] > t)

        if border_pos > 0:
            sky_b[x] = border_pos

    return sky_b

# Optimize energy function jn.
def CalculateOptimalBorder(image, thresh_min=5, thresh_max=600, search_step=5):
    grad = ColorToGradient(image)

    n = ((thresh_max - thresh_min) // search_step) + 1

    b_opt = None
    jn_max = 0

    for k in range(1, n + 1):
        t = thresh_min + ((thresh_max - thresh_min) // n - 1) * (k - 1)

        b_tmp = CalcBorder(grad, t)
        jn = EnergyFunction(b_tmp, image)

        if jn > jn_max:
            jn_max = jn
            b_opt = b_tmp

    return b_opt

# Called to check if only partial sky is detected.
def PartialSky(bopt, thresh4):
    return np.any(np.diff(bopt) > thresh4)

# Refinement function taken from Journal 2
def RefineSkyImage(bopt, image):
    sky_mask = GenerateMask(bopt, image)

    ground = np.ma.array(image, mask=cv.cvtColor(cv.bitwise_not(sky_mask), cv.COLOR_GRAY2BGR)).compressed()

    sky = np.ma.array(image, mask=cv.cvtColor(sky_mask, cv.COLOR_GRAY2BGR)).compressed()

    ground.shape = (ground.size//3, 3)
    sky.shape = (sky.size//3, 3)

    ret, label, center = cv.kmeans(np.float32(sky), 2, None, (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0), 10, cv.KMEANS_RANDOM_CENTERS)

    sigma_s1, mu_s1 = cv.calcCovarMatrix(sky[label.ravel() == 0], None, cv.COVAR_NORMAL | cv.COVAR_ROWS | cv.COVAR_SCALE)
    ic_s1 = cv.invert(sigma_s1, cv.DECOMP_SVD)[1]

    sigma_s2, mu_s2 = cv.calcCovarMatrix(sky[label.ravel() == 1], None, cv.COVAR_NORMAL | cv.COVAR_ROWS | cv.COVAR_SCALE)
    ic_s2 = cv.invert(sigma_s2, cv.DECOMP_SVD)[1]

    sigma_g, mu_g = cv.calcCovarMatrix(ground, None, cv.COVAR_NORMAL | cv.COVAR_ROWS | cv.COVAR_SCALE)

    icg = cv.invert(sigma_g, cv.DECOMP_SVD)[1]

    if cv.Mahalanobis(mu_s1, mu_g, ic_s1) > cv.Mahalanobis(mu_s2, mu_g, ic_s2):
        mu_s = mu_s1
        sigma_s = sigma_s1
        ics = ic_s1

    else:
        mu_s = mu_s2
        sigma_s = sigma_s2
        ics = ic_s2

    for x in range(image.shape[1]):
        cnt = np.sum(np.less(spatial.distance.cdist(image[0:bopt[x], x], mu_s, 'mahalanobis', VI=ics), spatial.distance.cdist(image[0:bopt[x], x], mu_g, 'mahalanobis', VI=icg)))

        if cnt < (bopt[x] / 2):
            bopt[x] = 0

    return bopt

# Main Function to decide our optimal sky detection function
def SkyDetect(enhanced_img, img):
    optimal_border_img = CalculateOptimalBorder(enhanced_img)

    if NoSkyRegion(optimal_border_img, enhanced_img.shape[0]/30, enhanced_img.shape[0]/4, 5):
        print("No sky detected")
        return

    elif PartialSky(optimal_border_img, enhanced_img.shape[1]/3):
        refined_optimal_border = RefineSkyImage(optimal_border_img, enhanced_img)
        DisplayMask(refined_optimal_border, img)

    else:
        DisplayMask(optimal_border_img, img)

# Called to check if the image contains no sky region based on the optimal border algorithm.
def NoSkyRegion(bopt, thresh1, thresh2, thresh3):
    border_ave = np.average(bopt)
    abs_sum_abs_diff = np.average(np.absolute(np.diff(bopt)))

    return border_ave < thresh1 or (border_ave < thresh2 and abs_sum_abs_diff > thresh3)

def GammaCorrection(img, gamma):
    invGamma = 1 / gamma

    table = [((i / 255) ** invGamma) * 255 for i in range(256)]
    table = np.array(table, np.uint8)

    return cv.LUT(img, table)

def AdjustBrightnessContrast(img, contrast:float=1.0, brightness:int=0):
    brightness += int(round(255*(1-contrast)/2))
    return cv.addWeighted(img, contrast, img, 0, brightness)

def ChangeSaturation(img, saturation):
    hsvImg = cv.cvtColor(img, cv.COLOR_BGR2HSV)
    hsvImg[...,1] = hsvImg[...,1] * saturation
    return cv.cvtColor(hsvImg, cv.COLOR_HSV2BGR)

# Inspired by Journal 1. Image is enhanced by
# directly extracting color channels that highlight the sky region.
# This function also applies the morphological operations described in
# the journal.
def ImageEnhancement(img):
    lab= cv.cvtColor(img, cv.COLOR_BGR2LAB)
    l_channel, a_channel, b_channel = cv.split(lab)

    # Instead of extracting the blue and violet colour channels, the brightness channel is normalised using a histogram.
    clahe = cv.createCLAHE(clipLimit=4.0, tileGridSize=(13,13))
    cl = clahe.apply(l_channel)

    blue_img = cv.merge((cl, a_channel, b_channel))

    enhanced_img = cv.cvtColor(blue_img, cv.COLOR_LAB2BGR)

    enhanced_img = ChangeSaturation(enhanced_img, 0.7)
    enhanced_img = AdjustBrightnessContrast(enhanced_img, 1.2, 1.2)
    enhanced_img = GammaCorrection(enhanced_img, 0.7)
    #enhanced_img = cv.fastNlMeansDenoisingColored(enhanced_img, None, 10, 10, 7, 15)

    # Dilation & Erosion
    kernel = np.ones((7, 7), np.uint8)
    enhanced_img = cv.dilate(enhanced_img, kernel, iterations=2)
    enhanced_img = cv.erode(enhanced_img, kernel, iterations=1)

    #CompareImgs(img, enhanced_img)

    return enhanced_img

# Driver Function
def main():
    dir = input("Enter Folder Name (Must be in imgs folder): ")
    output_dir = "output"

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    ext = ["png", "jpg", "jpeg"]

    files = []
    [files.extend(glob.glob("imgs/" + dir + "/*." + e)) for e in ext]
    images = [cv.imread(file) for file in files]

    for i in images:
        global ITER
        ITER = ITER + 1
        SkyDetect(ImageEnhancement(i), i)

if __name__ == '__main__':
    main()


import numpy as np
import network as nt
from network import UP, ACTIONS

X = []
model = nt.DQN.load_dqn("200 000_ite")

def FlappyPolicy(state, screen):
    s = nt.process_screen(screen)
    if len(X) < 4:
        X.append(s)
        return UP
    else:
        X.append(s)
        X.pop(0)
        x = np.stack(X, axis=-1)
        action = ACTIONS[nt.greedy_action(model, x)]
        return action


import numpy as np
import cv2

'''
cv2.Sobel(src, ddepth, dx, dy,
        dst=None, ksize= None, scale= None, delta= None,
        border_Type=None)

src: input image
ddepth: the same datatype if -1, recommends using cv2.CV_32F for more precise computation
dx : differential coef along the x axis
dy : differential coef along the y axis;
   : recommends using either dx = 1, dy =0 or dx = 0, dy = 1
dst: destination image
ksize: default = 3 (recommended)
scale: value to multiply to the 'raw' Sobal output
delta: bias to add to the 'raw' Sobal output --> to have a better look.
borderType: default= cv2.BORDER_DEFAULT (REFLECT_101)
'''
src = cv2.imread('./data/lenna.bmp', cv2.IMREAD_GRAYSCALE)
sobel = cv2.Sobel(src, ddepth= -1, dx=1, dy=0, delta= 128).astype(np.uint8)
# we use, delta = 128 to prevent saturation where pixels show gray to black patterns.
# In addition, black curves (lines) in the sobel filtered image indicates
# where the pixel values dramatically shifted.

dx = cv2.Sobel(src, ddepth= cv2.CV_32F, dx=1, dy=0, delta= 0)
dy = cv2.Sobel(src, ddepth= cv2.CV_32F, dx=0, dy=1, delta = 0)

# grad = dx, dy
mag = np.clip(cv2.magnitude(dx, dy), 0, 255).astype(np.uint8)

cv2.imshow('src', src)
cv2.imshow('sobel', sobel) # when dst dtype is the same with src dtype.
cv2.imshow('dx', dx)
cv2.imshow('dy', dy)
cv2.imshow('magnitude of gradient', mag)

cv2.waitKey()
cv2.destroyAllWindows()


import requests

url = "https://asr.iiit.ac.in/ssmtapi//"

payload={'lang': 'tel'}
files=[
  ('uploaded_file',('blob.wav',open('/home/raghava/Desktop/chat-ui-react/backend/blob.wav','rb'),'audio/wav'))
]
headers = {
  'accept': 'application/json'
}

response = requests.request("POST", url, headers=headers, data=payload, files=files, verify=False)


"""Mini-project 2 task 3 subtask e"""

def compose_archives(zip_one:str, zip_two:str, zip_result:str) -> None:
    """compose two archives into one
    >>> compose_archives("test_folder.zip", "just_folder.zip", "test_folder_result.zip")
    """
    import zipfile
    with zipfile.ZipFile(zip_one, 'r') as zip_file_one, zipfile.ZipFile(zip_two, 'r') as zip_file_two, zipfile.ZipFile(zip_result, 'w') as zip_file_result:
        for file in zip_file_one.namelist():
            zip_file_result.write(file)
        for file in zip_file_two.namelist():
            zip_file_result.write(file)

if __name__ == '__main__':
    # import doctest
    # print(doctest.testmod())
    import argparse
    parser = argparse.ArgumentParser(description="Program that makes zip file from 2 zip files")
    parser.add_argument('zip_one', help='First zip file')
    parser.add_argument('zip_two', help='Second zip file')
    parser.add_argument('zip_result', help='Resulting zip file')
    args = parser.parse_args()
    compose_archives(args.zip_one, args.zip_two, args.zip_result)


"""
problem tier : Gold 2 (solved.ac)
"""

import heapq
import sys
# sys.stdin = open('./input.txt', 'r')
input = sys.stdin.readline

N = int(input())

points = []
for i in range(N):
    h, o = map(int, input().split())
    if h > o:
        h, o = o, h
    points.append([h, o])

D = int(input())

points.sort(key=lambda x: x[1])

Q = []
count = 0
for p in points:
    heapq.heappush(Q, p[0])
    while len(Q) > 0 and Q[0] < p[1] - D:
        heapq.heappop(Q)

    count = max(count, len(Q))

print(count)


"""
Premium Question
"""
__author__ = 'Daniel'

class Solution(object):
    def validWordAbbreviation(self, word, abbr):
        """
        pointers
        :type word: str
        :type abbr: str
        :rtype: bool
        """
        w = 0
        a = 0
        while w < len(word) and a < len(abbr):
            if abbr[a].isdigit() and abbr[a] != '0':
                e = a
                while e < len(abbr) and abbr[e].isdigit(): e += 1
                num = int(abbr[a:e])
                a = e
                w += num
            else:
                if word[w] != abbr[a]:
                    return False

                w += 1
                a += 1

        return w == len(word) and a == len(abbr)

if __name__ == "__main__":
    assert Solution().validWordAbbreviation("internationalization", "i12iz4n") == True
    assert Solution().validWordAbbreviation("apple", "a2e") == False


num1 , num2 = int(input('Enter your first number : ')) , int(input('Enter your second number : '))
LCM =0
while num1 <= 0 or num2 <= 0 :
    num1 , num2 = int(input('Error! Enter positive number for your first number')) , int(input('Error! Enter positive number for your second number'))

max_num = max(num1, num2)
for i in range(max_num, num1 * num2 + 1):
    if (i % num1 == 0) and (i % num2 == 0):
        LCM = i
        break
print('Least common multiple : ', LCM)


from threading import Timer

import dash
import dash_table
import dash_core_components as dcc
import dash_html_components as html
import eventlet
import pandas as pd
from dash.dependencies import Input, Output
from flask import Flask
from flask_socketio import SocketIO

import sockettest

eventlet.monkey_patch()
server = Flask(__name__, static_url_path='')
socket = SocketIO(server, async_mode='eventlet')
app = dash.Dash(__name__, server=server)

app.scripts.config.serve_locally = True
app.css.config.serve_locally = True

df = pd.DataFrame({
    "Fruit": ["Apples", "Oranges", "Bananas", "Apples", "Oranges", "Bananas"],
    "Amount": [4, 1, 2, 2, 4, 5],
    "City": ["SF", "SF", "SF", "Montreal", "Montreal", "Montreal"]
})

def generate_table(dataframe, max_rows=100):
    return dash_table.DataTable(id='dt',
                                columns=[{"name": i, "id": i}
                                         for i in dataframe.columns],
                                data=dataframe.to_dict('records'),
                                page_size=30)

app.layout = html.Div(children=[
    sockettest.sockettest(),
    html.H1(children='MQTT Logger'),
    generate_table(df)
])

def change(id, val):
    socket.emit('call', {'id': id, 'val': val})

def timerCallback():
    global df
    Timer(1, timerCallback).start()
    change('dt', {'data': df.to_dict('records')})

if __name__ == '__main__':
    timerCallback()
    socket.run(server, debug=True, port=5000, host='0.0.0.0')


# 게임 개발

# 맵 크기 설정
n, m = map(int, input().split())

# 방문 위치 기억
d = [[0] * m for _ in range(n)]

# 캐릭터 위치, 방향 설정
x, y, direction = map(int, input().split())

# 시작 위치 카운트
d[x][y] = 1

# 게임 맵 생성
array = []
for i in range(n):
    array.append(list(map(int, input().split())))

# 이동방향 설정
# 0=북, 1=동, 2=남, 3=서
dx = [-1, 0, 1, 0]
dy = [0, 1, 0, -1]

# 왼쪽으로 회전하는 함수
def turn_left():
    global direction
    direction -= 1
    if direction == -1:
        direction = 3

# 시뮬레이션 시작
count = 1
turn_time = 0
while True:
    # 왼쪽으로 회전
    turn_left()
    nx = x + dx[direction]
    ny = y + dy[direction]
    # 회전한 이후 정면에 가보지 않은 육지인 칸이 존재하는 경우 이동
    if d[nx][ny] == 0 and array[nx][ny] == 0:
        d[nx][ny] = 1
        x = nx
        y = ny
        count += 1
        turn_time = 0
        continue
    # 회전한 이후 정면에 가보지 않은 칸이 없거나 바다인 경우
    else:
        turn_time += 1

    # 네 방향 모두 갈 수 없는 경우
    if turn_time == 4:
        # 뒤로 가는 좌표
        nx = x - dx[direction]
        ny = y - dy[direction]
        # 뒤로 갈 수 있는 경우
        if array[nx][ny] == 0:
            x = nx
            y = ny
        else:
            break
        turn_time = 0
print(count)


##python code for load and visualize a Linear regression type of dataset

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn import datasets
import matplotlib.pyplot as plt

##Generates sample data
X, y = datasets.make_regression(n_samples=100, n_features=1, noise=20, random_state=4)
##Split data in train/test samples/labels
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=1234)

##Plotting the dataset
fig = plt.figure(figsize=(8,6))
plt.scatter( X[:, 0], y, color = 'b', marker = 'o', s = 30)
plt.show()

##Showing data properties
print(f'The shape of X is {X_train.shape}.')
print(f'The shape of y is {y_train.shape}.')

##Testing the Linear Regression implementation
from ML_algorithms_scratch_implementation import LinearRegression

##Creates a regressor instance from LR class
regressor = LinearRegression(lr=0.01)
##Train the model
regressor.fit(X_train, y_train)
##Test the model
predicted_values = regressor.predict(X_test)

##Define the Mean Square Error for evaluate the model
def MSE(y_true, y_predicted):
    return np.mean((y_true - y_predicted)**2)

mse_value = MSE(y_test, predicted_values)

print(f'The MSE of the model is about {mse_value}.')


import random
import unittest
import io
from contextlib import redirect_stdout
import re

class testAutograde(unittest.TestCase):
    def test_1(self):
        goofy = DisneyBody('Goofy', ['wallet', 'paper', 'rock', 'scissors'])
        attribute_dict = vars(goofy)
        for key in attribute_dict:
            if isinstance(getattr(goofy, key), list):  # damn that's clever
                retrieved_pocket = getattr(goofy, key)
                pocket_key = key
            else:
                retrieved_name = attribute_dict[key]
                name_key = key
        self.assertTrue(getattr(goofy, name_key) == 'Goofy' and
                        getattr(goofy, pocket_key) == ['wallet', 'paper', 'rock', 'scissors'])

    def test_2(self):
        goofy = DisneyBody('Goofy', ['wallet', 'paper', 'rock', 'scissors'])
        goofy.add_content('meatball')
        attribute_dict = vars(goofy)
        for key in attribute_dict:
            if isinstance(getattr(goofy, key), list):
                retrieved_pocket = getattr(goofy, key)
                pocket_key = key
            else:
                retrieved_name = attribute_dict[key]
                name_key = key
        self.assertEqual(['wallet', 'paper', 'rock', 'scissors', 'meatball'], getattr(goofy, pocket_key),
                         "Added 'meatball' to Goofy with ['wallet', 'paper', 'rock', 'scissors'] "
                         "
Expect to get Goofy with ['wallet', 'paper', 'rock', 'scissors', 'meatball']"
                         "Goofy pocket is " + str(goofy.pocket))

    def test_3(self):
        goofy = DisneyBody('Goofy', ['wallet', 'paper', 'rock', 'scissors'])
        dr_wu = DisneyBody('Dr. Wu', ['strawberries'])
        goofy.add_content(dr_wu)
        attribute_dict = vars(goofy)
        for key in attribute_dict:
            if isinstance(getattr(goofy, key), list):
                retrieved_pocket = getattr(goofy, key)
                pocket_key = key
            else:
                retrieved_name = attribute_dict[key]
                name_key = key
        self.assertEqual(['wallet', 'paper', 'rock', 'scissors'], getattr(goofy, pocket_key),
                         "Added Dr. Wu (of type DisneyBody) to Goofy with ['strawberries']
"
                         "Expect to get Goofy with ['wallet', 'paper', 'rock', 'scissors'] (i.e. no change"
                         "to Goofy, Goofy does'nt care about strawberries.
"
                         "Goofy pocket is " + str(goofy.pocket))

    def test_4(self):
        goofy = DisneyBody('Goofy', ['wallet', 'paper', 'rock', 'scissors'])
        dr_wu = DisneyBody('Dr. Wu', ['pancreas', 'liver', 'heart', 'strawberries'])
        goofy.add_content(dr_wu)
        attribute_dict = vars(goofy)
        for key in attribute_dict:
            if isinstance(getattr(goofy, key), list):
                retrieved_pocket = getattr(goofy, key)
                pocket_key = key
            else:
                retrieved_name = attribute_dict[key]
                name_key = key

        self.assertEqual(['wallet', 'paper', 'rock', 'scissors', 'pancreas', 'liver', 'heart'],
                         getattr(goofy, pocket_key),
                         "Added Dr. Wu (of type DisneyBody) to Goofy with "
                         "['pancreas', 'liver', 'heart', 'strawberries']
"
                         "Expect to get Goofy with ['wallet', 'paper', 'rock', 'scissors', 'pancreas', 'liver', "
                         "'heart']] (i.e. Goofy does'nt care about strawberries but does care about organs.
"
                         "Goofy pocket is " + str(goofy.pocket)
                         )

    def test_5(self):
        from contextlib import redirect_stdout

        goofy = DisneyBody('Goofy', ['wallet', 'paper', 'rock', 'scissors'])

        f = io.StringIO()
        with redirect_stdout(f):
            print(goofy)
        print_string = f.getvalue()
        match_1 = len(re.findall(r"wallet", print_string, re.X | re.M | re.S))
        match_2 = len(re.findall(r"paper", print_string, re.X | re.M | re.S))
        match_3 = len(re.findall(r"rock", print_string, re.X | re.M | re.S))
        match_4 = len(re.findall(r"scissors", print_string, re.X | re.M | re.S))
        self.assertTrue(match_1 > 0 and match_2 > 0 and match_3 > 0 and match_4 > 0,
                        "Expect the __str__ method to return Goofy's pocket contents."
                        "Goofy's pocket is  ['wallet', 'paper', 'rock', 'scissors']
so return "
                        "should include wallet, paper, rock, scissors.
"
                        "Found this many matches:
"
                        "Wallet: " + str(match_1) + "
" + "paper: " + str(match_2) +
                        "
 rock: " + str(match_3) + "
"
                                                       "scissors: " + str(
                            match_4) + "
 in this output: " + print_string + "
")

if __name__ == '__main__':
    unittest.main()


import logging

l = logging.getLogger("archinfo.arch_amd64")

try:
    import capstone as _capstone
except ImportError:
    _capstone = None

try:
    import keystone as _keystone
except ImportError:
    _keystone = None

try:
    import unicorn as _unicorn
except ImportError:
    _unicorn = None

try:
    import pyvex as _pyvex
except ImportError:
    _pyvex = None

from .arch import Arch, register_arch, Endness, Register
from .tls import TLSArchInfo
from .archerror import ArchError

_NATIVE_FUNCTION_PROLOGS = [
        br"\x55\x48\x89\xe5", # push rbp; mov rbp, rsp
        br"\x48[\x83,\x81]\xec[\x00-\xff]", # sub rsp, xxx
    ]
# every function prolog can potentially be prefixed with endbr64
_endbr64 = b"\xf3\x0f\x1e\xfa"
_prefixed = [(_endbr64 + prolog) for prolog in _NATIVE_FUNCTION_PROLOGS]
_FUNCTION_PROLOGS = _prefixed + _NATIVE_FUNCTION_PROLOGS

class ArchAMD64(Arch):
    def __init__(self, endness=Endness.LE):
        if endness != Endness.LE:
            raise ArchError('Arch AMD64 must be little endian')
        super(ArchAMD64, self).__init__(endness)
        self.argument_register_positions = {
            self.registers['rdi'][0]: 0,
            self.registers['rsi'][0]: 1,
            self.registers['rdx'][0]: 2,
            self.registers['rcx'][0]: 3,  # Used for user calls
            self.registers['r10'][0]: 3,  # Used for Linux kernel calls
            self.registers['r8'][0]: 4,
            self.registers['r9'][0]: 5,
            # fp registers
            self.registers['xmm0'][0]: 0,
            self.registers['xmm1'][0]: 1,
            self.registers['xmm2'][0]: 2,
            self.registers['xmm3'][0]: 3,
            self.registers['xmm4'][0]: 4,
            self.registers['xmm5'][0]: 5,
            self.registers['xmm6'][0]: 6,
            self.registers['xmm7'][0]: 7
        } if _pyvex is not None else None
        if _unicorn:
            self.unicorn_flag_register = _unicorn.x86_const.UC_X86_REG_EFLAGS

    @property
    def capstone_x86_syntax(self):
        """
        The current syntax Capstone uses for x64. It can be 'intel' or 'at&t'
        """
        return self._cs_x86_syntax

    @capstone_x86_syntax.setter
    def capstone_x86_syntax(self, new_syntax):
        if new_syntax not in ('intel', 'at&t'):
            raise ArchError('Unsupported Capstone x86 syntax. It must be either "intel" or "at&t".')

        if new_syntax != self._cs_x86_syntax:
            self._cs = None
            self._cs_x86_syntax = new_syntax

    def _configure_capstone(self):
        self._cs.syntax = _capstone.CS_OPT_SYNTAX_ATT if self._cs_x86_syntax == 'at&t' else _capstone.CS_OPT_SYNTAX_INTEL

    @property
    def keystone_x86_syntax(self):
        """
        The current syntax Keystone uses for x86. It can be 'intel',
        'at&t', 'nasm', 'masm', 'gas' or 'radix16'
        """
        return self._ks_x86_syntax

    @keystone_x86_syntax.setter
    def keystone_x86_syntax(self, new_syntax):
        if new_syntax not in ('intel', 'at&t', 'nasm', 'masm', 'gas', 'radix16'):
            raise ArchError('Unsupported Keystone x86 syntax. It must be one of the following: '
                            '"intel", "at&t", "nasm", "masm", "gas" or "radix16".')

        if new_syntax != self._ks_x86_syntax:
            self._ks = None
            self._ks_x86_syntax = new_syntax

    def _configure_keystone(self):
        if self._ks_x86_syntax == 'at&t':
            self._ks.syntax = _keystone.KS_OPT_SYNTAX_ATT
        elif self._ks_x86_syntax == 'nasm':
            self._ks.syntax = _keystone.KS_OPT_SYNTAX_NASM
        elif self._ks_x86_syntax == 'masm':
            self._ks.syntax = _keystone.KS_OPT_SYNTAX_MASM
        elif self._ks_x86_syntax == 'gas':
            self._ks.syntax = _keystone.KS_OPT_SYNTAX_GAS
        elif self._ks_x86_syntax == 'radix16':
            self._ks.syntax = _keystone.KS_OPT_SYNTAX_RADIX16
        else:
            self._ks.syntax = _keystone.KS_OPT_SYNTAX_INTEL

    bits = 64
    vex_arch = "VexArchAMD64"
    vex_endness = "VexEndnessLE"
    name = "AMD64"
    qemu_name = 'x86_64'
    ida_processor = 'metapc'
    linux_name = 'x86_64'
    triplet = 'x86_64-linux-gnu'
    max_inst_bytes = 15
    ret_offset = 16
    vex_conditional_helpers = True
    syscall_num_offset = 16
    call_pushes_ret = True
    stack_change = -8
    initial_sp = 0x7ffffffffff0000
    call_sp_fix = -8
    memory_endness = Endness.LE
    register_endness = Endness.LE
    sizeof = {'short': 16, 'int': 32, 'long': 64, 'long long': 64}
    if _capstone:
        cs_arch = _capstone.CS_ARCH_X86
        cs_mode = _capstone.CS_MODE_64 + _capstone.CS_MODE_LITTLE_ENDIAN
    _cs_x86_syntax = None # Set it to 'att' in order to use AT&T syntax for x86
    if _keystone:
        ks_arch = _keystone.KS_ARCH_X86
        ks_mode = _keystone.KS_MODE_64 + _keystone.KS_MODE_LITTLE_ENDIAN
    _ks_x86_syntax = None
    uc_arch = _unicorn.UC_ARCH_X86 if _unicorn else None
    uc_mode = (_unicorn.UC_MODE_64 + _unicorn.UC_MODE_LITTLE_ENDIAN) if _unicorn else None
    uc_const = _unicorn.x86_const if _unicorn else None
    uc_prefix = "UC_X86_" if _unicorn else None
    function_prologs = _FUNCTION_PROLOGS
    function_epilogs = {
        br"\xc9\xc3", # leaveq; retq
        br"([^\x41][\x50-\x5f]{1}|\x41[\x50-\x5f])\xc3", # pop <reg>; retq
        br"\x48[\x83,\x81]\xc4([\x00-\xff]{1}|[\x00-\xff]{4})\xc3", #  add rsp, <siz>; retq
    }
    ret_instruction = b"\xc3"
    nop_instruction = b"\x90"
    instruction_alignment = 1
    register_list = [
        Register(name='rax', size=8, subregisters=[('eax', 0, 4),
                                                   ('ax', 0, 2),
                                                   ('al', 0, 1),
                                                   ('ah', 1, 1)],
                 general_purpose=True, linux_entry_value=0x1C),
        Register(name='rcx', size=8, subregisters=[('ecx', 0, 4),
                                                   ('cx', 0, 2),
                                                   ('cl', 0, 1),
                                                   ('ch', 1, 1)],
                 general_purpose=True, argument=True),
        Register(name='rdx', size=8, subregisters=[('edx', 0, 4),
                                                   ('dx', 0, 2),
                                                   ('dl', 0, 1),
                                                   ('dh', 1, 1)],
                 general_purpose=True, argument=True, linux_entry_value='ld_destructor'),
        Register(name='rbx', size=8, subregisters=[('ebx', 0, 4),
                                                   ('bx', 0, 2),
                                                   ('bl', 0, 1),
                                                   ('bh', 1, 1)],
                 general_purpose=True,
                 linux_entry_value=0),
        Register(name='rsp', size=8, subregisters=[('esp', 0, 4)], alias_names=('sp',),
                 general_purpose=True, default_value=(initial_sp, True, 'global')),
        Register(name='rbp', size=8, subregisters=[('ebp', 0, 4), ('bpl', 0, 1), ('bph', 1, 1)], alias_names=('bp',),
                 general_purpose=True, linux_entry_value=0),
        Register(name='rsi', size=8, subregisters=[('esi', 0, 4),
                                                   ('si', 0, 2),
                                                   ('sil', 0, 1),
                                                   ('sih', 1, 1)],
                 general_purpose=True, argument=True, linux_entry_value='argv'),
        Register(name='rdi', size=8, subregisters=[('edi', 0, 4),
                                                   ('di', 0, 2),
                                                   ('dil', 0, 1),
                                                   ('dih', 1, 1)],
                 general_purpose=True, argument=True, linux_entry_value='argc'),
        Register(name='r8', size=8, subregisters=[('r8d', 0, 4),
                                                  ('r8w', 0, 2),
                                                  ('r8b', 0, 1)],
                 general_purpose=True, argument=True),
        Register(name='r9', size=8, subregisters=[('r9d', 0, 4),
                                                  ('r9w', 0, 2),
                                                  ('r9b', 0, 1)],
                 general_purpose=True, argument=True),
        Register(name='r10', size=8, subregisters=[('r10d', 0, 4),
                                                   ('r10w', 0, 2),
                                                   ('r10b', 0, 1)],
                 general_purpose=True, argument=True),
        Register(name='r11', size=8, subregisters=[('r11d', 0, 4),
                                                   ('r11w', 0, 2),
                                                   ('r11b', 0, 1)],
                 general_purpose=True),
        Register(name='r12', size=8, subregisters=[('r12d', 0, 4),
                                                   ('r12w', 0, 2),
                                                   ('r12b', 0, 1)],
                 general_purpose=True,
                 linux_entry_value=0),
        Register(name='r13', size=8, subregisters=[('r13d', 0, 4),
                                                   ('r13w', 0, 2),
                                                   ('r13b', 0, 1)],
                 general_purpose=True,
                 linux_entry_value=0),
        Register(name='r14', size=8, subregisters=[('r14d', 0, 4),
                                                   ('r14w', 0, 2),
                                                   ('r14b', 0, 1)],
                general_purpose=True,
                 linux_entry_value=0),
        Register(name='r15', size=8, subregisters=[('r15d', 0, 4),
                                                   ('r15w', 0, 2),
                                                   ('r15b', 0, 1)],
                general_purpose=True,
                 linux_entry_value=0),
        Register(name='cc_op', size=8, default_value=(0, False, None), concrete=False, artificial=True),
        Register(name='cc_dep1', size=8, concrete=False, artificial=True),
        Register(name='cc_dep2', size=8, concrete=False, artificial=True),
        Register(name='cc_ndep', size=8, concrete=False, artificial=True, linux_entry_value=0),
        Register(name='d', size=8, alias_names=('dflag',),
                 default_value=(1, False, None), concrete=False),
        Register(name='rip', size=8, alias_names=('ip', 'pc'),
                 general_purpose=True),
        Register(name='ac', size=8, alias_names=('acflag',),concrete=False),
        Register(name='id', size=8, alias_names=('idflag',)),
        Register(name='fs', size=8, alias_names=('fs_const',),
                 default_value=(0x9000000000000000, True, 'global'),concrete=False),
        Register(name='sseround', size=8, vector=True,
                 default_value=(0, False, None)),
        Register(name='cr0', size=8),
        Register(name='cr2', size=8),
        Register(name='cr3', size=8),
        Register(name='cr4', size=8),
        Register(name='cr8', size=8),
        Register(name='ymm0', size=32, subregisters=[('xmm0', 0, 16),
                                                     ('xmm0lq', 0, 8),
                                                     ('xmm0hq', 8, 8),
                                                     ('ymm0hx', 16, 16)],
                 vector=True),
        Register(name='ymm1', size=32, subregisters=[('xmm1', 0, 16),
                                                     ('xmm1lq', 0, 8),
                                                     ('xmm1hq', 8, 8),
                                                     ('ymm1hx', 16, 16)],
                 vector=True),
        Register(name='ymm2', size=32, subregisters=[('xmm2', 0, 16),
                                                     ('xmm2lq', 0, 8),
                                                     ('xmm2hq', 8, 8),
                                                     ('ymm2hx', 16, 16)],
                 vector=True),
        Register(name='ymm3', size=32, subregisters=[('xmm3', 0, 16),
                                                     ('xmm3lq', 0, 8),
                                                     ('xmm3hq', 8, 8),
                                                     ('ymm3hx', 16, 16)],
                 vector=True),
        Register(name='ymm4', size=32, subregisters=[('xmm4', 0, 16),
                                                     ('xmm4lq', 0, 8),
                                                     ('xmm4hq', 8, 8),
                                                     ('ymm4hx', 16, 16)],
                 vector=True),
        Register(name='ymm5', size=32, subregisters=[('xmm5', 0, 16),
                                                     ('xmm5lq', 0, 8),
                                                     ('xmm5hq', 8, 8),
                                                     ('ymm5hx', 16, 16)],
                 vector=True),
        Register(name='ymm6', size=32, subregisters=[('xmm6', 0, 16),
                                                     ('xmm6lq', 0, 8),
                                                     ('xmm6hq', 8, 8),
                                                     ('ymm6hx', 16, 16)],
                 vector=True),
        Register(name='ymm7', size=32, subregisters=[('xmm7', 0, 16),
                                                     ('xmm7lq', 0, 8),
                                                     ('xmm7hq', 8, 8),
                                                     ('ymm7hx', 16, 16)],
                 vector=True),
        Register(name='ymm8', size=32, subregisters=[('xmm8', 0, 16),
                                                     ('xmm8lq', 0, 8),
                                                     ('xmm8hq', 8, 8),
                                                     ('ymm8hx', 16, 16)],
                 vector=True),
        Register(name='ymm9', size=32, subregisters=[('xmm9', 0, 16),
                                                     ('xmm9lq', 0, 8),
                                                     ('xmm9hq', 8, 8),
                                                     ('ymm9hx', 16, 16)],
                 vector=True),
        Register(name='ymm10', size=32, subregisters=[('xmm10', 0, 16),
                                                      ('xmm10lq', 0, 8),
                                                      ('xmm10hq', 8, 8),
                                                      ('ymm10hx', 16, 16)],
                 vector=True),
        Register(name='ymm11', size=32, subregisters=[('xmm11', 0, 16),
                                                      ('xmm11lq', 0, 8),
                                                      ('xmm11hq', 8, 8),
                                                      ('ymm11hx', 16, 16)],
                 vector=True),
        Register(name='ymm12', size=32, subregisters=[('xmm12', 0, 16),
                                                      ('xmm12lq', 0, 8),
                                                      ('xmm12hq', 8, 8),
                                                      ('ymm12hx', 16, 16)],
                 vector=True),
        Register(name='ymm13', size=32, subregisters=[('xmm13', 0, 16),
                                                      ('xmm13lq', 0, 8),
                                                      ('xmm13hq', 8, 8),
                                                      ('ymm13hx', 16, 16)],
                 vector=True),
        Register(name='ymm14', size=32, subregisters=[('xmm14', 0, 16),
                                                      ('xmm14lq', 0, 8),
                                                      ('xmm14hq', 8, 8),
                                                      ('ymm14hx', 16, 16)],
                 vector=True),
        Register(name='ymm15', size=32, subregisters=[('xmm15', 0, 16),
                                                      ('xmm15lq', 0, 8),
                                                      ('xmm15hq', 8, 8),
                                                      ('ymm15hx', 16, 16)],
                 vector=True),
        Register(name='ftop', size=4, floating_point=True, default_value=(0, False, None), artificial=True),
        Register(name='fpreg', size=64, subregisters=[('mm0', 0, 8),
                                                      ('mm1', 8, 8),
                                                      ('mm2', 16, 8),
                                                      ('mm3', 24, 8),
                                                      ('mm4', 32, 8),
                                                      ('mm5', 40, 8),
                                                      ('mm6', 48, 8),
                                                      ('mm7', 56, 8)],
                 alias_names=('fpu_regs',), floating_point=True),
        Register(name='fptag', size=8, alias_names=('fpu_tags',),
                 floating_point=True, default_value=(0, False, None)),
        Register(name='fpround', size=8, floating_point=True, default_value=(0, False, None)),
        Register(name='fc3210', size=8, floating_point=True),
        Register(name='emnote', size=4, artificial=True),
        Register(name='cmstart', size=8),
        Register(name='cmlen', size=8),
        Register(name='nraddr', size=8),
        Register(name='gs', size=8, alias_names=('gs_const',), concrete=False),
        Register(name='ip_at_syscall', size=8, concrete=False, artificial=True),
    ]

    symbol_type_translation = {
        10: 'STT_GNU_IFUNC',
        'STT_LOOS': 'STT_GNU_IFUNC'
    }
    got_section_name = '.got.plt'
    ld_linux_name = 'ld-linux-x86-64.so.2'
    elf_tls = TLSArchInfo(2, 704, [16], [8], [0], 0, 0)

register_arch([r'.*amd64|.*x64|.*x86_64|.*metapc'], 64, Endness.LE, ArchAMD64)


import re

def parse_input(filename):
    with open(filename) as f:
        parts = f.read().split("\n\n")

        matches = re.findall(r"([a-z ]+): ([0-9]+)\-([0-9]+) or ([0-9]+)\-([0-9]+)", parts[0])
        rules = {name: [range(int(a0), int(a1)+1), range(int(b0), int(b1)+1)] for (name, a0, a1, b0, b1) in matches}

        my_ticket = list(map(int, parts[1].split("\n")[1].split(",")))

        tickets = [list(map(int, line.strip().split(","))) for line in parts[2].split("\n")[1:]]

        return my_ticket, tickets, rules

def is_valid(ticket, rules):
    ranges = [r for value in rules.values() for r in value]
    return all(map(lambda t: any(map(lambda r: t in r, ranges)), ticket))

def solve_positions(matches):
    if len(matches) == 0:
        return {}

    s = { k: v[0] for (k, v) in matches.items() if len(v) == 1}

    for k in s.keys():
        del matches[k]

    for key, value in matches.items():
        matches[key] = [v for v in value if v not in s.values()]

    return s | solve_positions(matches)

def main(filename):
    my_ticket, tickets, rules = parse_input(filename)

    valid_tickets = [t for t in tickets if is_valid(t, rules)] + [my_ticket]
    ticket_length = len(valid_tickets[0])

    matches = {}

    for name, ranges in rules.items():
        m = []
        for i in range(ticket_length):
            if all([any(map(lambda r: t[i] in r, ranges)) for t in valid_tickets]):
                m.append(i)
        matches[name] = m

    positions = solve_positions(matches)
    res = 1
    for s, i in positions.items():
        if "departure" in s:
            res *= my_ticket[i]
    return res

result = main("input.txt")
print(result)


mylist=['x','a','m','a','x']
start = 0
end = len(mylist)-1
count=0
for item in mylist:
    if mylist[start] == mylist[end]:
        start+=1
        end-=1
        count=1
    else:
        start += 1
        end -= 1
        count-=1

if count > 0:
    print(f'Yes! Given list {mylist} is  a Palindrome !')
else:
    print(f'No! Given list {mylist} is not a Palindrome !')


import warnings
warnings.filterwarnings('ignore')
import pandas
import datetime
import numpy
from scipy import stats
import sqlite3
import os
import sys
module_path = os.path.abspath(os.path.join('../data_pipeline/'))
if module_path not in sys.path:
    sys.path.append(module_path)
import omn_utils

# set some variables
# omn db
omnDBDir = "../data/sqlite3/"
omnDbName = "omni_sw_imf.sqlite"
omnTabName = "imf_sw"
omnParams = ["By", "Bz", "Bx", "Vx", "Np"]
omnTimeDelay = 10 #minutes
marchTime = 30 # minutes
onsetTimeInterval = 15 #minutes
# onset data
smlFname="../data/filtered-20190103-22-53-substorms.csv"
# some plotting vars
predDateRange = [datetime.datetime(1996,1,1), datetime.datetime(2018,1,1)]#None

# load onset data
ssOnsetDF = pandas.read_csv(smlFname,\
                                parse_dates=["Date_UTC"])
ssOnsetDF.columns = [ "date", "mlat", "mlt" ]
print("loaded onset data")
# get the time range and load OMNI data
if predDateRange is not None:
    omnStartDate = predDateRange[0] - datetime.timedelta(minutes=60)
    omnEndDate = predDateRange[1]
else:
    omnStartDate = ssOnsetDF["date"].min() - datetime.timedelta(minutes=60)
    omnEndDate = ssOnsetDF["date"].max()
# create the obj and load data
omnObj = omn_utils.OmnData(omnStartDate, omnEndDate, omnDBDir,\
                   omnDbName, omnTabName,\
                   True, None,\
                    imf_normalize=False,\
                    db_time_resolution=1,\
                    omn_train_params = omnParams)
# add a 10 minute time delay to the omni dataset
omnObj.omnDF["datetime"] = omnObj.omnDF["datetime"] + datetime.timedelta(minutes=omnTimeDelay)
omnObj.omnDF.set_index(omnObj.omnDF["datetime"], inplace=True)

# loop through and make the predictions
# using the conditions
predDict = {}
predDict["date"] = []
predDict["prediction"] = []
predDict["actual"] = []
predDict["onsetDate"] = []
predDict["triggerTime"] = []
nUnknown = 0
nFalse = 0
nTrue = 0

# Loop through the daterange and get number of onsets/false onsets etc
loopDate = predDateRange[0]
currMonth = loopDate.month
currYear = loopDate.year
print("working through month, year-->",currMonth, currYear)
while loopDate <= predDateRange[1]:
    if (currMonth != loopDate.month) | (currYear != loopDate.year):
        currMonth = loopDate.month
        currYear = loopDate.year
        print("working through month, year-->",currMonth, currYear)
    _sTime = loopDate - datetime.timedelta(minutes=30)
    _eTime = loopDate
    _currOmn = omnObj.omnDF[ _sTime : _eTime ].dropna()# Move to the next date
    loopDate += datetime.timedelta(minutes=marchTime)
    # Rule 1 : in the previous 30 minutes, atleast 22 minutes should have
    #           negative Bz. For this we also need to check if we have 30
    #           values of Bz (sometimes data is missing)
    # there should be 30 values
    if _currOmn.shape[0] < 30:
        print("insufficient data", _currOmn.shape)
        predDict["date"].append( loopDate )
        predDict["prediction"].append( "U" )
        predDict["triggerTime"].append( datetime.datetime(1970,1,1) )
        nUnknown += 1
        # check if there is an onset in the next 10 minutes
        _chkOnsets = ssOnsetDF[ (ssOnsetDF["date"] >= loopDate) &\
                            (ssOnsetDF["date"] <= (loopDate+datetime.timedelta(minutes=onsetTimeInterval))) ]
        if _chkOnsets.shape[0] > 0:
            predDict["onsetDate"].append( _chkOnsets["date"].tolist()[0] )
            predDict["actual"].append( "T" )
        else:
            predDict["onsetDate"].append( -1. )
            predDict["actual"].append( "F" )
        continue
    else:
        # count number of negative Bz
        nBz = _currOmn[_currOmn["Bz"] < 0.].shape[0]
        if nBz >= 22:
            # criteria1 succesful!
            predDict["date"].append( loopDate )
            predDict["prediction"].append( "T" )
            predDict["triggerTime"].append( loopDate )
            nTrue += 1
            # check if there is an onset in the next 10 minutes
            _chkOnsets = ssOnsetDF[ (ssOnsetDF["date"] >= loopDate) &\
                                (ssOnsetDF["date"] <= (loopDate+datetime.timedelta(minutes=onsetTimeInterval))) ]
            if _chkOnsets.shape[0] > 0:
                predDict["onsetDate"].append( _chkOnsets["date"].tolist()[0] )
                predDict["actual"].append( "T" )
            else:
                predDict["onsetDate"].append( -1. )
                predDict["actual"].append( "F" )
        else:
            # criteria1 failed! No SS
            predDict["date"].append( loopDate )
            predDict["prediction"].append( "F" )
            predDict["triggerTime"].append( datetime.datetime(1980,1,1) )
            nFalse += 1
            # check if there is an onset in the next 10 minutes
            _chkOnsets = ssOnsetDF[ (ssOnsetDF["date"] >= loopDate) &\
                                (ssOnsetDF["date"] <= (loopDate+datetime.timedelta(minutes=onsetTimeInterval))) ]
            if _chkOnsets.shape[0] > 0:
                predDict["onsetDate"].append( _chkOnsets["date"].tolist()[0] )
                predDict["actual"].append( "T" )
            else:
                predDict["onsetDate"].append( -1. )
                predDict["actual"].append( "F" )
            continue

print("nTrue-->", nTrue)
print("nFalse-->", nFalse)
print("nUnknown-->", nUnknown)
# finally calculate all the TPs/FPs etc
predDF = pandas.DataFrame(predDict)

# save this DF
predDF.to_csv("../data/lyons_marching_criteria1.csv")

print("##############TP count##############")
print(predDF[ (predDF["prediction"] == "T") & (predDF["actual"] == "T") ].shape)
print("##############TP count##############")

print("##############FP count##############")
print(predDF[ (predDF["prediction"] == "T") & (predDF["actual"] == "F") ].shape)
print("##############FP count##############")

print("##############TN count##############")
print(predDF[ (predDF["prediction"] == "F") & (predDF["actual"] == "F") ].shape)
print("##############TN count##############")

print("##############FN count##############")
print(predDF[ (predDF["prediction"] == "F") & (predDF["actual"] == "T") ].shape)
print("##############FN count##############")


"""Объявите класс с именем DataBase (база данных - БД), объекты которого создаются командой:

db = DataBase(path)
где path - путь к файлу с данными БД (строка).

Также в классе DataBase нужно объявить следующие методы:

write(self, record) - для добавления новой записи в БД, представленной объектом record;
read(self, pk) - чтение записи из БД (возвращает объект Record) по ее уникальному идентификатору pk
(уникальное целое положительное число); запись ищется в значениях словаря (см. ниже)

Каждая запись БД должна описываться классом Record, а объекты этого класса создаваться командой:

record = Record(fio, descr, old)

где fio - ФИО некоторого человека (строка); descr - характеристика человека (строка);
old - возраст человека (целое число).

В каждом объекте класса Record должны формироваться следующие локальные атрибуты:

pk - уникальный идентификатор записи (число: целое, положительное); формируется автоматически при создании
каждого нового объекта;
fio - ФИО человека (строка);
descr - характеристика человека (строка);
old - возраст человека (целое число).

Реализовать для объектов класса Record вычисление хэша по атрибутам: fio и old (без учета регистра).
Если они одинаковы для разных записей, то и хэши должны получаться равными. Также для объектов класса Record
с одинаковыми хэшами оператор == должен выдавать значение True, а с разными хэшами - False.

Хранить записи в БД следует в виде словаря dict_db (атрибут объекта db класса DataBase), ключами которого являются
объекты класса Record, а значениями список из объектов с равными хэшами:

dict_db[rec1] = [rec1, rec2, ..., recN]

где rec1, rec2, ..., recN - объекты класса Record с одинаковыми хэшами.

Для наполнения БД прочитайте строки из входного потока с помощью команды:

lst_in = list(map(str.strip, sys.stdin.readlines()))
где каждая строка представлена в формате:

"ФИО; характеристика; возраст"

Например:

Балакирев С.М.; программист; 33
Кузнецов А.В.; разведчик-нелегал; 35
Суворов А.В.; полководец; 42
Иванов И.И.; фигурант всех подобных списков; 26
Балакирев С.М.; преподаватель; 37

Каждая строка должна быть представлена объектом класса Record и записана в БД db (в словарь db.dict_db).

P.S. На экран ничего выводить не нужно."""

import sys

class DataBase:
    def __init__(self, path):
        self.path = path
        self.dict_db = {}

    def write(self, record):
        if record in self.dict_db:
            self.dict_db[record].append(record)
        else:
            self.dict_db[record] = [record]

    def read(self, pk):
        for lst_val in self.dict_db.values():
            for item in lst_val:
                if item.pk == pk:
                    return item

class Record:
    _counter = 0

    def __new__(cls, *args, **kwargs):
        obj = super().__new__(cls)
        setattr(obj, 'pk', cls._counter)
        cls._counter += 1
        return obj

    def __init__(self, fio, descr, old):
        self.fio = fio
        self.descr = descr
        self.old = old

    def __hash__(self):
        return hash((self.fio.lower(), self.old))

    def __eq__(self, other):
        return hash(self) == hash(other)

db = DataBase('www.vk.com')

lst_in = list(map(str.strip, sys.stdin.readlines()))
lst_4_obj = [x.split('; ') for x in lst_in]
for item in lst_4_obj:
    db.write(Record(*item))

db22345 = DataBase('123')
r1 = Record('fio', 'descr', 10)
r2 = Record('fio', 'descr', 10)
assert r1.pk != r2.pk, "равные значения атрибута pk у разных объектов класса Record"

db22345.write(r2)
r22 = db22345.read(r2.pk)
assert r22.pk == r2.pk and r22.fio == r2.fio and r22.descr == r2.descr and r22.old == r2.old, "при операциях write и read прочитанный объект имеет неверные значения атрибутов"

assert len(db22345.dict_db) == 1, "неверное число объектов в словаре dict_db"

fio = lst_in[0].split(';')[0].strip()
v = list(db.dict_db.values())
if fio == "Балакирев С.М.":
    assert len(v[0]) == 2 and len(v[1]) == 1 and len(v[2]) == 1 and len(
        v[3]) == 1, "неверно сформирован словарь dict_db"

if fio == "Гейтс Б.":
    assert len(v[0]) == 2 and len(v[1]) == 2 and len(v[2]) == 1 and len(
        v[3]) == 1, "неверно сформирован словарь dict_db"


import cv2
import numpy as np
import pandas as pd
import pathlib

import tensorflow as tf
import random
from backend.models import CNNModel,ResBlock

# Loading Models

# Emotion Model

emotion_model = tf.keras.models.load_model('backend/models/vgg.h5')

# Face Model

cascade_path = pathlib.Path(cv2.__file__).parent.absolute() / "data/haarcascade_frontalface_default.xml"
face_model = cv2.CascadeClassifier(str(cascade_path))

# categories

categories = ['angry','disgust','fear','happy','neutral','sad','surprise']

# music
df_music = pd.read_csv('backend/dataframes/music.csv')

class ImageAndMusic():

    def __init__(self,image):
        self.face_model = face_model
        self.emotion_model = emotion_model
        self.image = image

    # predict image

    def predict_image_and_recommend(self):

        gray=cv2.cvtColor(self.image,cv2.COLOR_BGR2GRAY)
        faces = face_model.detectMultiScale(gray,
            scaleFactor = 1.1,
            minNeighbors = 5,
            minSize = (30,30),
            flags = cv2.CASCADE_SCALE_IMAGE
            )
        if len(faces) != 0 :
            for (x, y, width, height ) in faces:
                x_min,y_min,x_max,y_max = x , y, x + width, y + height
                prediction = self.predict(self.image[y_min:y_max , x_min:x_max])
        else:
            # Or Predict Image
            prediction = self.predict(self.image)

        recommended_song = self.recommend(prediction)
        return prediction,recommended_song

    # Predict image

    def predict(self,image):
        # Resize the image to the desired size (48x48 pixels with 3 channels)

        desired_size = (48, 48)
        resized_img = cv2.resize(image, desired_size)

        prediction = emotion_model.predict(resized_img.reshape(1, 48, 48, 3))
        # prediction = emotion_model.predict(resized_img)

        final_prediction = np.argmax(prediction)
        return(categories[final_prediction].upper())

    def recommend(self,category):

        if category in ['fear','angry','surprise']:
            df_filtered = df_music[df_music["mood"] == 'Energetic'].reset_index(drop=True)
            recommended_song = random.choice(df_filtered['name'])
        elif category == 'happy':
            df_filtered = df_music[df_music["mood"] == 'Happy'].reset_index(drop=True)
            recommended_song = random.choice(df_filtered['name'])
        elif category == 'sad':
            df_filtered = df_music[df_music["mood"] == 'Sad'].reset_index(drop=True)
            recommended_song = random.choice(df_filtered['name'])
        else:
            df_filtered = df_music[df_music["mood"] == 'Calm'].reset_index(drop=True)
            recommended_song = random.choice(df_filtered['name'])

        return recommended_song


from typing import Dict, List

import pytest
from factory.base import FactoryMetaClass

@pytest.fixture
def check_product_list(product_factory: FactoryMetaClass) -> List[Dict]:
    product1 = product_factory()
    product2 = product_factory()
    return [
        {
            "id": product1.pk,
            "title": product1.title,
            "category": {"id": product1.category.pk, "title": product1.category.title},
        },
        {
            "id": product2.pk,
            "title": product2.title,
            "category": {"id": product2.category.pk, "title": product2.category.title},
        },
    ]

@pytest.fixture
def check_product_retrieve(product_factory: FactoryMetaClass) -> Dict:
    product = product_factory()
    return {
        "id": product.pk,
        "title": product.title,
        "category": {"id": product.category.pk, "title": product.category.title},
        "institutions": [],
    }

@pytest.fixture
def check_category_list(category_factory: FactoryMetaClass) -> List[Dict]:
    category1 = category_factory()
    category2 = category_factory()
    return [{"id": category1.pk, "title": category1.title}, {"id": category2.pk, "title": category2.title}]


import pygame
from pygame.sprite import Sprite

class Star(Sprite):
    """A class to represent a single star in the sky."""

    def __init__(self, ai_game):
        """Initialize the sky and set its starting position."""
        super().__init__()
        self.screen = ai_game.screen

        # Load the star image and set its rect attribute.
        self.image = pygame.image.load('img/star_xs.png')
        self.rect = self.image.get_rect()

        # Start each new star near  the top left of the screen.
        self.rect.x = self.rect.width
        self.rect.y = self.rect.height

        # Store the star's exact horizontal and vertical positions.
        self.x = float(self.rect.x)
        self.y = float(self.rect.y)


from lib.mqttclient import MQTTClient
import ubinascii
import machine
import log
import relay
import time
import uasyncio
import utils

topics = [ b'sprinkler/zone1',
           b'sprinkler/zone2',
           b'sprinkler/zone3',
           b'sprinkler/zone4',
           b'reset',
           b'update',
           b'schedule']
#mqtt_server = '98.51.182.241'

r = relay.Relay()

def restart_and_reconnect():
    log.info('Failed toconnect to MQTT broker, Reconnecting...')
    time.sleep(10)
    machine.reset()

def initPubSub():
    try:
        client = connect_and_subscribe()
    except OSError as e:
        restart_and_reconnect()
    while True:
        try:
            client.check_msg()
            await uasyncio.sleep_ms(500)
            #time.sleep (0.5)
            r.setBoardLed1(1)
            await uasyncio.sleep_ms(500)
#            time.sleep(0.5)
            r.setBoardLed1(0)
        except OSError as e:
            restart_and_reconnect()

def sub_cb(topic, msg):
    log.info((topic, msg))
    log.info('ESP received on message {}'.format(msg))

    if topic == b'sprinkler/zone1' and msg == b'on':
        r.setRelay1(0)
        # create a class for pin to off and on
    elif topic == b'sprinkler/zone1' and msg == b'off':
        r.setRelay1(1)
        # create a class for pin to off and on
    elif topic == b'sprinkler/zone2' and msg == b'on':
        r.setRelay2(0)
        # create a class for pin to off and on
    elif topic == b'sprinkler/zone2' and msg == b'off':
        r.setRelay2(1)
        # create a class for pin to off and on
    elif topic == b'sprinkler/zone3' and msg == b'on':
        r.setRelay3(0)
        # create a class for pin to off and on
    elif topic == b'sprinkler/zone3' and msg == b'off':
        r.setRelay3(1)
        # create a class for pin to off and on
    elif topic == b'sprinkler/zone4' and msg == b'on':
        r.setRelay4(0)
        # create a class for pin to off and on
    elif topic == b'sprinkler/zone4' and msg == b'off':
        r.setRelay4(1)
        # create a class for pin to off and on
    elif topic == b'reset':
        machine.reset()
        #reboot esp
    elif topic == b'update':
        log.info('ESP received update')
        #update code OTA
        try:
          from lib.ota_updater import OTAUpdater
          otaUpdater = OTAUpdater('https://github.com/aparajita-gupta/sprinkler-system', github_src_dir='sprinklerSystem', main_dir='esp32', secrets_file="secret.json")
          if otaUpdater.install_update_if_available():
            log.info("Updated successfully")
          else:
            log.info("Already Updated")
        except:
          log.info("Update Failed...")
    elif topic == b'schedule':
        log.info('ESP received schedule')
        try:
            f = open('./app/schedule.json', 'w')
            f.write(msg)
            f.close()
            machine.reset()
        except:
            log.info("Couldn't update schedule")
    else:
        log.info('ESP could not understand message')

def connect_and_subscribe():
    mqtt_server = utils.getMQTTServerIP()
    client_id = ubinascii.hexlify(machine.unique_id())
    global client
    log.info('Connecting server: {}'.format(mqtt_server))
    client = MQTTClient(client_id, mqtt_server)
    log.info('Connected to %s MQTT broker'.format(mqtt_server))
    client.set_callback(sub_cb)
    client.connect()
    for topic_sub in topics:
        client.subscribe(topic_sub)
        log.info('Subscribed to %s topic' % (topic_sub))
    # We may not need return calue client because it is global
    return client


#%%
import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate
import pandas as pd
import pickle

n_i0=1
n_e0=1
Length=8000
Ifr=8001
Rmn=3600
Ifr_pos=10
Te0=1

xj_1=np.linspace(0,Ifr_pos,Ifr)
xj_2=np.linspace(Ifr_pos,100,Rmn+1)
xi_1=np.empty(Ifr-1)
xi_2=np.empty(Rmn)
for i in range(Ifr-1):
    xi_1[i]=0.5*(xj_1[i]+xj_1[i+1])
for i in range(Rmn):
    xi_2[i]=0.5*(xj_2[i]+xj_2[i+1])
xj=np.append(xj_1,xj_2[1:])
xi=np.append(xi_1,xi_2)

#initial ion and potential distribution
ni0_i=np.heaviside(-(xi-Ifr_pos),1)
ni0_j=np.heaviside(-(xj-Ifr_pos),1)
Phi0_1=np.zeros(Ifr-1)+1
Phi0_2=np.zeros(Rmn)+1
Phi0j_1=np.zeros(Ifr)+1
Phi0j_2=np.zeros(Rmn+1)+1

#Matrix EQN on leftside of ion front
def Matrix1(xi_l,xi_r,Phi1,phi2,T):

    x_temp=np.append(-xi_l[0],xi_l)
    x=np.append(x_temp,xi_r[0])
    Phi=np.append(Phi1,phi2)
    A=np.zeros((Ifr+1,Ifr+1))

    A[0][0] = -1/(x[1]-x[0])
    A[0][1] = 1/(x[1]-x[0])

     #last row
    f1 = 1/(4*np.sqrt(2*T))*((np.exp(Phi[-1]/T))**(1/2)+(np.exp(Phi[-2]/T))**(1/2))
    g1 = (np.sqrt(T/2))*((np.exp(Phi[-1]/T))**(1/2)+(np.exp(Phi[-2]/T))**(1/2))
    A[-1][-2] = -(1-f1*(x[-1]-x[-2]))
    A[-1][-1] = 1+f1*(x[-1]-x[-2])

    for i in range(1,Ifr):
        A[i][i-1] = 2/((x[i+1]-x[i-1])*(x[i]-x[i-1]))
        A[i][i] = -(2/((x[i+1]-x[i])*(x[i]-x[i-1])))-(np.exp(Phi[i-1]/T)/T)
        A[i][i+1] = 2/((x[i+1]-x[i-1])*(x[i+1]-x[i]))

    return(A)

def VecB1(xi_l,xi_r,Phi1,phi2, ni,T):
    B=np.zeros(Ifr+1)
    B[0]=0

    Phi=np.append(Phi1,phi2)
    for i in range(1,Ifr):
        B[i]=np.exp(Phi1[i-1]/T)*(1-(Phi1[i-1]/T))-ni[i-1]

     #boundary conditions
    f1 = 1/(4*np.sqrt(2*T))*((np.exp(Phi[-1]/T))**(1/2)+(np.exp(Phi[-2]/T))**(1/2))
    g1 =(np.sqrt(T/2))*((np.exp(Phi[-1]/T))**(1/2)+(np.exp(Phi[-2]/T))**(1/2))

    B[-1] = (f1*(Phi[-1]+Phi[-2])-g1)*(xi_r[0]-xi_l[-1])
    return(B)

def potential1(xi_l,xi_r,Phi1,phi2, ni,T):
    Phi1 = np.linalg.solve(Matrix1(xi_l,xi_r,Phi1,phi2,T),VecB1(xi_l,xi_r,Phi1,phi2, ni,T))
    return(Phi1)

def PrecisePhi1(xi_l,xi_r,Phi1,phi2, ni,T):
    Phi_new=np.zeros((9,Ifr+1))
    Phi_new[0] = np.linalg.solve(Matrix1(xi_l,xi_r,Phi1,phi2,T),VecB1(xi_l,xi_r,Phi1,phi2, ni,T))

    for i in range(8):
        Phi_new[i+1] = potential1(xi_l,xi_r,Phi_new[i][1:-1],Phi_new[i][-1], ni,T)
    return(Phi_new[-1])

Phi_1=PrecisePhi1(xi_1,xi_2,Phi0_1,Phi0_2[0],ni0_i,Te0)[1:]

def PrecisePhi2(Phi1,Phi2,xi,xj,phi_bdr,T):
    Phi=np.zeros(Rmn+1)
    x=np.append(xi,xj[-1])
    Phi[0]=Phi1[-1]
    for i in range(1,Rmn+1):
        Phi[i]=Phi[0]-2*T*np.log(np.exp(Phi[0]/(2*T))*(x[i]-x[0])/np.sqrt(2*T)+1)
    return(Phi[:-1],Phi[-1])

Phi_2,phi_end=PrecisePhi2(Phi_1,Phi0_2,xi_2,xj_2,Phi0j_2[-1],Te0)
E_end=np.sqrt(2*Te0)*np.exp(phi_end/(2*Te0))

Phi=np.append(Phi_1,Phi_2[1:])
plt.plot(xi,Phi)
plt.show()

def field_fl(x_i1,x_i2, Phi1): #field on the left side of ion front
    N = Ifr
    x_i=np.append(x_i1,x_i2[0])
    E = np.zeros(N-1)
    E[0]=-0.5*(Phi1[1]-Phi1[0])/(x_i1[1]-x_i1[0])
    for i in range(1,N-1):
        E[i] = -0.5*((Phi1[i]-Phi1[i+1])/(x_i[i]-x_i[i+1])+(Phi1[i]-Phi1[i-1])/(x_i[i]-x_i[i-1]))
    return(E)
E_fl=field_fl(xi_1,xi_2,Phi_1)

def field_fr(x_i1,x_i2, Phi1,Phi2,T):
    N = Rmn

    E=np.sqrt(2*T)*np.exp(Phi2/(2*T))
    return(E)

E_fr=field_fr(xi_1,xi_2,Phi_1,Phi_2,Te0)

E=np.append(E_fl,E_fr)
plt.plot(xi,E)
plt.show()

def fieldj_fl(xj,ni,Phi,T):
    Ej=np.zeros(Ifr)
    for i in range(Ifr-1):
        Ej[i+1]=Ej[i]-(np.exp(Phi[i]/T)-ni[i])*(xj[i+1]-xj[i])
    return(Ej)

#%%
#Having done with the initial condition
#We now deal with time evolution of the system
def time1(Phi1,Phi2, xj_l,xj_r, xi_l,xi_r, E1,E2,n_ion0l,n_ion0r,T,E_last,Phi_last):
    #time span for each time steps
    dt = 0.025
    dT=0 #Temperature change
    N=Ifr
    M=Rmn

    Tt=np.zeros(5000) #Temperature with repect to time
    Tt[0]=T #1MeV
    Tt[1]=T

    enum=np.zeros((3000,2)) #Number of electrons, to monitor conservation of electrons

    #all the values are stored in matrices, where the values at time j*dt are in the jth column
    #First, initialize these arrays
    Phi1_all = np.zeros((N,5000)) #Array describing potential left to the ion-front on i-grid system
    Phi1_all[0][0]=Phi1[0]
    Phi1_all[:,0]=Phi1
    Phij_end=np.zeros(5000)

    Phi2_all=np.zeros((M,5000))    #Potential on grid points at the right side of ion-front on i-grid system
    Phi2_all[:,0]=Phi2
    Phij_end[0]=Phi_last

    xj_lall = np.zeros((N,5000))   #j-grids left to the ion-front
    xj_lall[0][0]=-xj_l[0]
    xj_lall[:,0] = xj_l

    xj_rall=np.zeros((M+1,5000))   #j-grids right to the ion-front
    xj_rall[:,0]=xj_r

    xi_lall=np.zeros((N-1,5000))   #i-grids left to the ion-front
    xi_lall[:,0]=xi_l
    xi_rall=np.zeros((M,5000))
    xi_rall[:,0]=xi_r

    ni_lall=np.zeros((N-1,5000))
    ni_lall[:,0]=n_ion0l

    vj_fl = np.zeros((N,3000))

    vi_fl=np.zeros((N-1,3000))

    Ej_fl=np.zeros((N,5000))

    Ei_fl=np.zeros((N-1,5000))
    Ei_fr=np.zeros((Rmn,5000))

    Ei_fl[:,0]=E1
    Ei_fr[:,0]=E2

    Ej_fl[:,0]=fieldj_fl(xj_l,n_ion0l,Phi1,Te0)

    #Time evolution with the leapfrog method
    for t in range(4):

        #Velocity
        for j in range(N):
            vj_fl[j][t+1]=vj_fl[j][t]+Ej_fl[j][t]*dt

        for i in range(N-1):
            vi_fl[i][t+1]=0.5*(vj_fl[i][t+1]+vj_fl[i+1][t+1])

        #j-grid position
        for j in range(N):
            xj_lall[j][t+1]=xj_lall[j][t]+vj_fl[j][t]*dt+0.5*Ej_fl[j][t]*dt*dt

        xj_rall[0][t+1]=xj_lall[-1][t+1]
        for j in range(1,M+1):
            xj_rall[j][t+1]=xj_rall[j][t]+xj_lall[-1][t+1]-xj_lall[-1][t]

        #i-grid position
        for i in range(N-1):
            xi_lall[i][t+1]=0.5*(xj_lall[i][t+1]+xj_lall[i+1][t+1])

        for i in range(M):
            xi_rall[i][t+1]=0.5*(xj_rall[i][t+1]+xj_rall[i+1][t+1])

        #ion density, on i-grid of course

        for i in range(N-1):
            ni_lall[i][t+1]=ni_lall[i][t]*((xj_lall[i+1][t]-xj_lall[i][t])/(xj_lall[i+1][t+1]-xj_lall[i][t+1]))

        tail=0
        for i in range(N-1):
            enum[t][0]=enum[t][0]+np.exp(Phi1_all[i][t]/Tt[t])*(xj_lall[i+1][t]-xj_lall[i][t])

        enum[t][0]=enum[t][0]+Ej_fl[-1][t] #Ej_fl[-1][t] happends to be equal to Ne after ion front quantitively

        if t>=1:

            energy1=2*Tt[t-1]*Ej_fl[-1][t-1]
            energy2=2*Tt[t]*Ej_fl[-1][t]

            xj_a=np.append(xj_lall[:,t-1],xj_rall[1:,t-1])
            xj_b=np.append(xj_lall[:,t],xj_rall[1:,t])
            Ej_a=Ej_fl[:,t-1]
            Ej_b=Ej_fl[:,t]
            vj_a=vj_fl[:,t-1]
            vj_b=vj_fl[:,t]
            na=ni_lall[:,t-1]
            nb=ni_lall[:,t]
            for i in range(N-1):
                energy1=energy1+((xj_a[i+1]-xj_a[i])/3)*(Ej_a[i+1]**2+(Ej_a[i]**2)+Ej_a[i+1]*Ej_a[i])
                energy2=energy2+((xj_b[i+1]-xj_b[i])/3)*(Ej_b[i+1]**2+(Ej_b[i]**2)+Ej_b[i+1]*Ej_b[i])

            for i in range(N-1):
                energy1=energy1+(na[i]/3)*(xj_a[i+1]-xj_a[i])*(vj_a[i+1]**2+(vj_a[i])**2+vj_a[i+1]*vj_a[i])
                energy2=energy2+(nb[i]/3)*(xj_b[i+1]-xj_b[i])*(vj_b[i+1]**2+(vj_b[i])**2+vj_b[i+1]*vj_b[i])

            dT=(1/xj_lall[-1,0])*(energy1-energy2)

            Tt[t+1]=Tt[t]+dT
            print("dT=",dT)
            enum[t][1]=enum[t][0]-enum[t-1][0]

        if Tt[t+1]<0:
            print("Error!T<0")
            break

        Phi1_all[:,t+1]=PrecisePhi1(xi_lall[:,t+1],xi_rall[:,t+1],Phi1_all[:-1,t],Phi2_all[0,t],ni_lall[:,t+1],Tt[t+1])[1:]
        Phi2_all[:,t+1],Phij_end[t+1]=PrecisePhi2(Phi1_all[:,t+1],Phi2_all[:,t],xi_rall[:,t+1],xj_rall[:,t+1],Phij_end[t],Tt[t+1])

        Ei_fl[:,t+1]=field_fl(xi_lall[:,t+1],xi_rall[:,t+1],Phi1_all[:,t+1])
        Ei_fr[:,t+1]=field_fr(xi_lall[:,t+1],xi_rall[:,t+1],Phi1_all[:,t+1],Phi2_all[:,t+1],Tt[t+1])

        Ej_fl[:,t+1]=fieldj_fl(xj_lall[:,t+1],ni_lall[:,t+1],Phi1_all[:,t+1],Tt[t+1])

        print("Number of Electrons:",enum[t])

        print("Next Temperature:", Tt[t+1])

        plt.plot(np.append(xi_lall[:,t+1],xi_rall[:,t+1]),np.append(Phi1_all[:-1,t+1],Phi2_all[:,t+1]))
        plt.title("Potential")
        plt.show()
        plt.plot(np.append(xi_lall[:,t+1],xi_rall[:,t+1]),np.append(Ei_fl[:,t+1],Ei_fr[:,t+1]))
        plt.title("Electric field")
        plt.show()
        plt.plot(np.append(xi_lall[:,t+1],xi_rall[:,t+1]),np.append(ni_lall[:,t+1],np.zeros(Rmn)),label='ni')
        ne=np.append(Phi1_all[:-1,t+1],Phi2_all[:,t+1])
        plt.plot(np.append(xi_lall[:,t+1],xi_rall[:,t+1]),np.exp(ne/Tt[t+1]),label='ne')
        plt.title("charge distribution")
        plt.legend()
        plt.show()

    return(Tt,Phi1_all, Phi2_all, Phij_end,xj_lall,xj_rall,xi_lall,xi_rall,ni_lall,vj_fl,Ej_fl,\
        Ei_fl,Ei_fr,enum)

T_t,Pot1,Pot2,Pot_end,jgrid_lall,jgrid_rall,igrid_lall,igrid_rall,ion_lall,velj_fl,\
    Elcj_fl,Elci_fl,Elci_fr,ENum=time1(Phi_1,Phi_2,xj_1,xj_2,xi_1,xi_2,E_fl,E_fr,ni0_i[:Ifr-1],ni0_i[Ifr-1:],Te0,E_end,phi_end)