import casadi as cd

# from sys import path
# path.append(r'/home/smarnlab/Code_base')
# import casadi as cd

import matplotlib.pyplot as plt
import numpy as np

### Four-Fingered Optimization

class GraspOptimSolver_4F:
    def __init__(self, solver_name="ipopt", gripper_type="FishGrip4F"):
        ### Parameters
        gripperConstraints = {
            'FishGri_4F':{'palm_dists':[30, 30, 30*np.sqrt(2), 30*np.sqrt(2)], 'dis_open':[-15, 30], 'cor_ang':0, 'cor_line_id':['14', '23']},
            'RochuGr_4F':{'palm_dists':[30, 30, 30*np.sqrt(2), 30*np.sqrt(2)], 'dis_open':[-15, 30], 'cor_ang':np.pi/2, 'cor_line_id':['13', '24']},
        }
        self.solver_name = solver_name
        self.gripper_type = gripper_type
        self.init_palm_dists = gripperConstraints[gripper_type]['palm_dists']
        self.init_dis_open = gripperConstraints[gripper_type]['dis_open']
        self.init_cor_ang = gripperConstraints[gripper_type]['cor_ang']
        self.init_cor_line_id = gripperConstraints[gripper_type]['cor_line_id']

    def refresh_mapsfuns(self, angle_map, certainty_map):
        ### Angle and prob maps
        xgrid_dense = np.linspace(0,300,300, endpoint=False)
        ygrid_dense = np.linspace(0,300,300, endpoint=False)
        self.M_desen_ang = np.copy(angle_map)
        self.M_desen_cer = np.copy(certainty_map)
        self.uncer_max = np.max(self.M_desen_cer)

        ang_flat = self.M_desen_ang.ravel(order='F')
        Rf_ang = cd.interpolant('Rf_ang','bspline',[xgrid_dense, ygrid_dense], ang_flat)
        cer_flat = self.M_desen_cer.ravel(order='F')
        Rf_cer = cd.interpolant('Rf_cer','bspline',[xgrid_dense, ygrid_dense], cer_flat)

        ### Initial cd Functions
        a = cd.MX.sym("a", 1, 1)
        b = cd.MX.sym("b", 1, 1)
        self.func_Rf_ang = cd.Function('func_ang', [a, b], [Rf_ang(cd.vcat([a, b]))], ['a', 'b'], ['Rf1'])
        self.func_Rf_cer = cd.Function('func_cer', [a, b], [Rf_cer(cd.vcat([a, b]))], ['a', 'b'], ['Rf2'])
        c = cd.MX.sym("c", 1, 1)
        self.angle_tocs = cd.Function('angle_tocs', [c], [cd.hcat([cd.cos(c), cd.sin(c)])])
        var_u1 = cd.MX.sym('var_u1')
        var_v1 = cd.MX.sym('var_v1')
        var_u2 = cd.MX.sym('var_u2')
        var_v2 = cd.MX.sym('var_v2')
        self.distance = cd.Function('distance', [var_u1, var_v1, var_u2, var_v2], [cd.sqrt((var_u1-var_u2)**2 + (var_v1-var_v2)**2)])
        ### Convinent functions
        var_x1 = cd.MX.sym('var_x1')
        var_y1 = cd.MX.sym('var_y1')
        var_x2 = cd.MX.sym('var_x2')
        var_y2 = cd.MX.sym('var_y2')
        self.AngleBT = cd.Function('AngleBT', [var_x1, var_y1, var_x2, var_y2], [cd.acos((var_x1*var_x2 + var_y1*var_y2)/(cd.norm_2(cd.hcat([var_x1, var_y1])) \
                * cd.norm_2(cd.hcat([var_x2, var_y2])))) * cd.sign(var_x1*var_y2 - var_x2*var_y1)], ['var_x1', 'var_y1', 'var_x2', 'var_y2'], ['AngleBT'])
        uu = cd.MX.sym("uu", 1, 1)
        vv = cd.MX.sym("vv", 1, 1)
        tt = cd.MX.sym("tt", 1, 1)
        self.fun_mapangdiff = cd.Function('fun_mapangdiff', [uu, vv, tt], \
                [0.5*cd.norm_2(self.angle_tocs(self.func_Rf_ang(uu, vv)) - self.angle_tocs(tt) + 1e-6)])
        tt1 = cd.MX.sym("tt1", 1, 1)
        tt2 = cd.MX.sym("tt2", 1, 1)
        self.fun_angdiff = cd.Function('fun_angdiff', [tt1, tt2], [0.5*cd.norm_2(self.angle_tocs(tt2) - self.angle_tocs(tt1) + 1e-6)])
        # self.fun_angdiff_crosssin = cd.Function('fun_angdiff', [tt1, tt2], [cd.fabs(self.angle_tocs(tt1)[0] *self.angle_tocs(tt2)[1] - self.angle_tocs(tt2)[0] *self.angle_tocs(tt1)[1])])

    def init_solver(self):
        ### init_solver
        self.opti_solver = cd.Opti()

        ### Variables
        self.ut1 = self.opti_solver.variable()
        self.vt1 = self.opti_solver.variable()
        self.th1 = self.opti_solver.variable()
        self.ut2 = self.opti_solver.variable()
        self.vt2 = self.opti_solver.variable()
        self.th2 = self.opti_solver.variable()
        self.ut3 = self.opti_solver.variable()
        self.vt3 = self.opti_solver.variable()
        self.th3 = self.opti_solver.variable()
        self.ut4 = self.opti_solver.variable()
        self.vt4 = self.opti_solver.variable()
        self.th4 = self.opti_solver.variable()
        self.dis_open = self.opti_solver.variable()

        self.up1 = self.ut1 - self.dis_open * cd.sin(self.th1)
        self.vp1 = self.vt1 + self.dis_open * cd.cos(self.th1)
        self.up2 = self.ut2 - self.dis_open * cd.sin(self.th2)
        self.vp2 = self.vt2 + self.dis_open * cd.cos(self.th2)
        self.up3 = self.ut3 - self.dis_open * cd.sin(self.th3)
        self.vp3 = self.vt3 + self.dis_open * cd.cos(self.th3)
        self.up4 = self.ut4 - self.dis_open * cd.sin(self.th4)
        self.vp4 = self.vt4 + self.dis_open * cd.cos(self.th4)

        ### Parameters
        self.unc_thread = self.opti_solver.parameter()
        self.opti_solver.set_value(self.unc_thread, 0.7) # to be changed
        self.palm_d12 = self.opti_solver.parameter()
        self.palm_d23 = self.opti_solver.parameter()
        self.palm_d13 = self.opti_solver.parameter()
        self.palm_d24 = self.opti_solver.parameter()
        self.opti_solver.set_value(self.palm_d12, self.init_palm_dists[0])
        self.opti_solver.set_value(self.palm_d23, self.init_palm_dists[1])
        self.opti_solver.set_value(self.palm_d13, self.init_palm_dists[2])
        self.opti_solver.set_value(self.palm_d24, self.init_palm_dists[3])
        self.dis_open_min = self.opti_solver.parameter()
        self.dis_open_max = self.opti_solver.parameter()
        self.opti_solver.set_value(self.dis_open_min, self.init_dis_open[0])
        self.opti_solver.set_value(self.dis_open_max, self.init_dis_open[1])

        ### Objective function
        u_tp1, v_tp1 = (self.ut1 - self.up1), (self.vt1 - self.vp1)
        u_tp2, v_tp2 = (self.ut2 - self.up2), (self.vt2 - self.vp2)
        u_p13, v_p13 = (self.up1 - self.up3), (self.vp1 - self.vp3)
        u_p14, v_p14 = (self.up1 - self.up4), (self.vp1 - self.vp4)
        u_p23, v_p23 = (self.up2 - self.up3), (self.vp2 - self.vp3)
        u_p24, v_p24 = (self.up2 - self.up4), (self.vp2 - self.vp4)
        vector_bank = {
            '13': [u_p13*cd.sign(self.dis_open), v_p13*cd.sign(self.dis_open)],
            '14': [u_p14*cd.sign(self.dis_open), v_p14*cd.sign(self.dis_open)],
            '23': [u_p23*cd.sign(self.dis_open), v_p23*cd.sign(self.dis_open)],
            '24': [u_p24*cd.sign(self.dis_open), v_p24*cd.sign(self.dis_open)],
        }
        self.ang_tp1p_a = self.AngleBT(-vector_bank[self.init_cor_line_id[0]][1], vector_bank[self.init_cor_line_id[0]][0], -v_tp1, u_tp1)
        self.ang_tp2p_b = self.AngleBT(-vector_bank[self.init_cor_line_id[1]][1], vector_bank[self.init_cor_line_id[1]][0], -v_tp2, u_tp2)

        obj_ang_corr_1 = self.fun_angdiff(self.ang_tp1p_a, 0)
        obj_ang_corr_2 = self.fun_angdiff(self.ang_tp2p_b, 0)

        obj_ang_anti_1 = cd.norm_2(self.angle_tocs(self.th1) + self.angle_tocs(self.th3))
        obj_ang_anti_2 = cd.norm_2(self.angle_tocs(self.th2) + self.angle_tocs(self.th4))
        # obj_ang_anti_1 = self.fun_angdiff(self.th1, self.th3 + cd.pi)
        # obj_ang_anti_2 = self.fun_angdiff(self.th2, self.th4 + cd.pi)
        # obj_ang_corr_3 = self.fun_angdiff(self.th1, self.th2 + self.init_cor_ang)
        # obj_ang_corr_4 = self.fun_angdiff(self.th3, self.th4 + self.init_cor_ang)

        self.obj_anb = (1.0/2.0)*(obj_ang_corr_1 + obj_ang_corr_2)
        # self.obj_ang = (1.0/4.0)*(obj_ang_anti_1 + obj_ang_anti_2 + obj_ang_corr_3 + obj_ang_corr_4)
        self.obj_ang = (1.0/2.0)*(obj_ang_anti_1 + obj_ang_anti_2)
        self.obj_cer = (1.0/4.0)*(self.func_Rf_cer(self.ut1, self.vt1) + self.func_Rf_cer(self.ut2, self.vt2)
            + self.func_Rf_cer(self.ut3, self.vt3) + self.func_Rf_cer(self.ut4, self.vt4))
        # self.objfun = cd.fabs(self.dis_open) + 100*(self.obj_ang + self.obj_anb + (1 - self.obj_cer))
        self.objfun = 100*(self.obj_anb + (1 - self.obj_cer))
        # self.objfun = 100*(self.obj_ang + self.obj_anb + (1 - self.obj_cer))
        self.opti_solver.minimize(self.objfun)

        ### Constraints
        # board constraints of image
        range_uv = 300
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.ut1, range_uv))
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.ut2, range_uv))
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.ut3, range_uv))
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.ut4, range_uv))
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.vt1, range_uv))
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.vt2, range_uv))
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.vt3, range_uv))
        self.opti_solver.subject_to(self.opti_solver.bounded(0, self.vt4, range_uv))
        # Distance constraints between palm point postitions

        # self.opti_solver.subject_to(self.distance(self.up1, self.vp1, self.up2, self.vp2) == self.palm_d12)
        # self.opti_solver.subject_to(self.distance(self.up2, self.vp2, self.up3, self.vp3) == self.palm_d23)
        # self.opti_solver.subject_to(self.distance(self.up3, self.vp3, self.up4, self.vp4) == self.palm_d12)
        # self.opti_solver.subject_to(self.distance(self.up4, self.vp4, self.up1, self.vp1) == self.palm_d23)
        # self.opti_solver.subject_to(self.distance(self.up1, self.vp1, self.up3, self.vp3) == self.palm_d13)
        # self.opti_solver.subject_to(self.distance(self.up2, self.vp2, self.up4, self.vp4) == self.palm_d24)

        dist_epsi_thresd = 1
        self.opti_solver.subject_to(cd.fabs(self.distance(self.up1, self.vp1, self.up2, self.vp2) - self.palm_d12) < dist_epsi_thresd)
        self.opti_solver.subject_to(cd.fabs(self.distance(self.up2, self.vp2, self.up3, self.vp3) - self.palm_d23) < dist_epsi_thresd)
        self.opti_solver.subject_to(cd.fabs(self.distance(self.up3, self.vp3, self.up4, self.vp4) - self.palm_d12) < dist_epsi_thresd)
        self.opti_solver.subject_to(cd.fabs(self.distance(self.up4, self.vp4, self.up1, self.vp1) - self.palm_d23) < dist_epsi_thresd)
        self.opti_solver.subject_to(cd.fabs(self.distance(self.up1, self.vp1, self.up3, self.vp3) - self.palm_d13) < dist_epsi_thresd)
        self.opti_solver.subject_to(cd.fabs(self.distance(self.up2, self.vp2, self.up4, self.vp4) - self.palm_d24) < dist_epsi_thresd)

        # Distance constraints between fingertips and palm points
        self.opti_solver.subject_to(self.opti_solver.bounded(self.dis_open_min, self.dis_open, self.dis_open_max))
        # threshold constraints from probability maps
        self.opti_solver.subject_to(self.func_Rf_cer(self.ut1, self.vt1) > self.unc_thread)
        self.opti_solver.subject_to(self.func_Rf_cer(self.ut2, self.vt2) > self.unc_thread)
        self.opti_solver.subject_to(self.func_Rf_cer(self.ut3, self.vt3) > self.unc_thread)
        self.opti_solver.subject_to(self.func_Rf_cer(self.ut4, self.vt4) > self.unc_thread)
        # Angle map - Angle match constraints
        # For consistency at pi and -pi, cos/sin operation are adopted

        self.opti_solver.subject_to(self.func_Rf_ang(self.ut1, self.vt1) == self.th1) # self.th1
        self.opti_solver.subject_to(self.func_Rf_ang(self.ut2, self.vt2) == self.th2) # self.th2
        self.opti_solver.subject_to(self.func_Rf_ang(self.ut3, self.vt3) == self.th3) # self.th3
        self.opti_solver.subject_to(self.func_Rf_ang(self.ut4, self.vt4) == self.th4) # self.th4

        cs_epsi_thresd = 5e-2 # 5e-2: 3 degs
        # vertical to each other
        self.opti_solver.subject_to(cd.fabs(self.angle_tocs(self.th1) + self.angle_tocs(self.th3)) < cs_epsi_thresd)
        self.opti_solver.subject_to(cd.fabs(self.angle_tocs(self.th2) + self.angle_tocs(self.th4)) < cs_epsi_thresd)

        # Finger orientation (to palm) constraints
        # cs_epsi_thresd = 5e-2
        # self.opti_solver.subject_to(self.fun_angdiff(self.th1, self.th3 + cd.pi) < cs_epsi_thresd)
        # self.opti_solver.subject_to(self.fun_angdiff(self.th2, self.th4 + cd.pi) < cs_epsi_thresd)
        # self.opti_solver.subject_to(self.fun_angdiff(self.th1, self.th2 + self.init_cor_ang) < cs_epsi_thresd)
        # self.opti_solver.subject_to(self.fun_angdiff(self.th3, self.th4 + self.init_cor_ang) < cs_epsi_thresd)

        # self.opti_solver.subject_to(self.fun_angdiff_crosssin(self.ang_tp1p14, 0) < cs_epsi_thresd)
        # self.opti_solver.subject_to(self.fun_angdiff_crosssin(self.ang_tp2p23, 0) < cs_epsi_thresd)
        incre_bias = cd.pi/180.0 * 3
        self.opti_solver.subject_to(self.opti_solver.bounded(- incre_bias, self.ang_tp1p_a, + incre_bias))
        self.opti_solver.subject_to(self.opti_solver.bounded(- incre_bias, self.ang_tp2p_b, + incre_bias))
        # self.opti_solver.subject_to(self.fun_angdiff(self.ang_tp1p14, 0) < cs_epsi_thresd)
        # self.opti_solver.subject_to(self.fun_angdiff(self.ang_tp2p23, 0) < cs_epsi_thresd)

    @staticmethod
    def plot_variable(ut1, vt1, th1, ut2, vt2, th2, ut3, vt3, th3, ut4, vt4, th4, opt_a, opti_b, opti_c, opti_d, opti_all):
        print("[{:.2f}] [{:.2f}] [{:.2f}] / [{:.2f}] [{:.2f}] [{:.2f}] / [{:.2f}] [{:.2f}] [{:.2f}] / [{:.2f}] [{:.2f}] [{:.2f}] opti: {:.2f}, {:.2f}, {:.2f}, {:.2f} == {:.2f}"\
                .format(ut1, vt1, th1, ut2, vt2, th2, ut3, vt3, th3, ut4, vt4, th4, opt_a, opti_b, opti_c, opti_d, opti_all))

    def iterate_solver(self, max_iter=300, total_intervals=5, trials_each_interval=20, rich_output=False):
        ### Solver
        s_opts = {
            # 'tol':10,
            # 'compl_inf_tol':100,
            # 'dual_inf_tol': 10,
            'max_iter': max_iter,

            'mu_target': 1e-1,
            'acceptable_tol':1, #50,
            'acceptable_iter':3,
            'acceptable_dual_inf_tol':1, #50,
            'acceptable_constr_viol_tol':1e-2,
            'acceptable_compl_inf_tol':1,
            'acceptable_obj_change_tol':1e-2,

            'print_level': 5 if rich_output else 1,
            'print_user_options':'yes' if rich_output else 'no',
            'print_options_documentation':'no',
            }

        self.opti_solver.solver(self.solver_name, {'ipopt':s_opts, 'print_time': 0})

        ### Debug settings output
        if rich_output:
            self.opti_solver.callback(lambda i: self.plot_variable(
                self.opti_solver.debug.value(self.ut1),
                self.opti_solver.debug.value(self.vt1),
                self.opti_solver.debug.value(self.th1),
                self.opti_solver.debug.value(self.ut2),
                self.opti_solver.debug.value(self.vt2),
                self.opti_solver.debug.value(self.th2),
                self.opti_solver.debug.value(self.ut3),
                self.opti_solver.debug.value(self.vt3),
                self.opti_solver.debug.value(self.th3),
                self.opti_solver.debug.value(self.ut4),
                self.opti_solver.debug.value(self.vt4),
                self.opti_solver.debug.value(self.th4),
                self.opti_solver.debug.value(self.dis_open),
                self.opti_solver.debug.value(self.obj_cer),
                self.opti_solver.debug.value(self.obj_ang),
                self.opti_solver.debug.value(self.obj_anb),
                self.opti_solver.debug.value(self.objfun),
                )
            )
        else: self.opti_solver.callback()

        ### Start Optimize
        solutions = []
        best_solution = {'objfun':1e+10, 'state': 'F'}
        solution_succ_num = 0

        for unc_interval_i in range(total_intervals):
            uncer_favor = self.uncer_max * float(total_intervals - unc_interval_i) / (total_intervals - unc_interval_i + 3)
            # interval_decay_rates = np.linspace(decrease_rate, 0.5, total_intervals, endpoint=False)
            # uncer_favor = self.uncer_max * interval_decay_rates[unc_interval_i]
            index_good = np.array(np.nonzero(self.M_desen_cer > uncer_favor)).T

            ### try select [trials_each_interval] times at unc_interval_i
            for trail_i in range(trials_each_interval):
                print('\n========= Trial_{} selection in unc_favor: {} ========='.format(trail_i, uncer_favor))
                point_inits, point_angles = [], np.zeros(4)
                for try_sample_i in range(20):
                    raint = np.random.randint(index_good.shape[0], size=4)
                    point_inits = index_good[raint]
                    for index in range(point_inits.shape[0]): point_angles[index] = self.func_Rf_ang(point_inits[index][0], point_inits[index][1])
                    var_angles = np.sqrt(np.mean(np.cos(point_angles))**2 + np.mean(np.sin(point_angles))**2)
                    if var_angles < 0.2: break
                    # else: print(point_angles, 'not good, resampling....')
                point_argsort = np.argsort(point_angles)
                index_good_sorted = point_inits[point_argsort[::-1]]

                # print('point_inits', point_inits, 'index_good_sorted', index_good_sorted)

                # index_turn = np.array([[0, 1, 2, 3], [1, 2, 3, 0], [2, 3, 0, 1], [3, 0, 1, 2]])
                index_turn = np.array([[0, 1, 2, 3], [1, 2, 3, 0]])
                for one_order in index_turn:
                    print('\n =={}== Trial_{} selection in unc_favor: {} ========= with order:'.format(solution_succ_num, trail_i, uncer_favor), one_order)
                    point1 = index_good_sorted[one_order[0]]
                    point2 = index_good_sorted[one_order[1]]
                    point3 = index_good_sorted[one_order[2]]
                    point4 = index_good_sorted[one_order[3]]
                    ang1 = self.func_Rf_ang(point1[0],point1[1]) + 1e-3 # for ERR gradient
                    ang2 = self.func_Rf_ang(point2[0],point2[1]) - 1e-3 # for ERR gradient
                    ang3 = self.func_Rf_ang(point3[0],point3[1]) + 1e-3 # for ERR gradient
                    ang4 = self.func_Rf_ang(point4[0],point4[1]) - 1e-3 # for ERR gradient

                    print("Try starting points: ", point1, point2, point3, point4, ang1, ang2, ang3, ang4)
                    self.opti_solver.set_initial(self.ut1, point1[0])
                    self.opti_solver.set_initial(self.vt1, point1[1])
                    self.opti_solver.set_initial(self.ut2, point2[0])
                    self.opti_solver.set_initial(self.vt2, point2[1])
                    self.opti_solver.set_initial(self.ut3, point3[0])
                    self.opti_solver.set_initial(self.vt3, point3[1])
                    self.opti_solver.set_initial(self.ut4, point4[0])
                    self.opti_solver.set_initial(self.vt4, point4[1])
                    self.opti_solver.set_initial(self.th1, ang1)
                    self.opti_solver.set_initial(self.th2, ang2)
                    self.opti_solver.set_initial(self.th3, ang3)
                    self.opti_solver.set_initial(self.th4, ang4)
                    init_dis_open_midr = 0.5*(self.init_dis_open[0] + np.random.rand(1)[0] * (self.init_dis_open[1] - self.init_dis_open[0]))
                    # (self.init_dis_open[0] + self.init_dis_open[1])*0.5 + 2.0
                    self.opti_solver.set_initial(self.dis_open, init_dis_open_midr)
                    self.opti_solver.set_value(self.unc_thread, 0.7 * uncer_favor)
                    # self.opti_solver.set_value(self.dis_min, dis_min_favor)
                    # self.opti_solver.set_value(self.dis_max, dis_max_favor)

                    try:
                        sol = self.opti_solver.solve()
                    except Exception as e:
                        print("Failure occurred at Trial {} selection in unc_favor: {} ".format(trail_i, uncer_favor))

                        # print("POSSIBLE infeasibility")
                        # self.opti_solver.debug.show_infeasibilities()

                        solution = {
                            'state':'F',
                            'iters': self.opti_solver.stats()['iter_count'],
                            'uncer_favor':uncer_favor,
                            'ut1':self.opti_solver.debug.value(self.ut1),
                            'ut2':self.opti_solver.debug.value(self.ut2),
                            'ut3':self.opti_solver.debug.value(self.ut3),
                            'ut4':self.opti_solver.debug.value(self.ut4),
                            'vt1':self.opti_solver.debug.value(self.vt1),
                            'vt2':self.opti_solver.debug.value(self.vt2),
                            'vt3':self.opti_solver.debug.value(self.vt3),
                            'vt4':self.opti_solver.debug.value(self.vt4),
                            'th1':self.opti_solver.debug.value(self.th1),
                            'th2':self.opti_solver.debug.value(self.th2),
                            'th3':self.opti_solver.debug.value(self.th3),
                            'th4':self.opti_solver.debug.value(self.th4),
                            'dis_open':self.opti_solver.debug.value(self.dis_open),
                            'obj_cer':self.opti_solver.debug.value(self.obj_cer),
                            'obj_ang':self.opti_solver.debug.value(self.obj_ang),
                            'obj_anb':self.opti_solver.debug.value(self.obj_anb),
                            'objfun':self.opti_solver.debug.value(self.objfun),
                        }
                    else:
                        print('\n\nSuccessing: in trial: {}\n\n'.format(trail_i))
                        solution = {
                            'state':'S',
                            'iters': self.opti_solver.stats()['iter_count'],
                            'uncer_favor':uncer_favor,
                            'ut1':sol.value(self.ut1),
                            'ut2':sol.value(self.ut2),
                            'ut3':sol.value(self.ut3),
                            'ut4':sol.value(self.ut4),
                            'vt1':sol.value(self.vt1),
                            'vt2':sol.value(self.vt2),
                            'vt3':sol.value(self.vt3),
                            'vt4':sol.value(self.vt4),
                            'th1':sol.value(self.th1),
                            'th2':sol.value(self.th2),
                            'th3':sol.value(self.th3),
                            'th4':sol.value(self.th4),
                            'dis_open':sol.value(self.dis_open),
                            'obj_cer':sol.value(self.obj_cer),
                            'obj_ang':sol.value(self.obj_ang),
                            'obj_anb':sol.value(self.obj_anb),
                            'objfun':sol.value(self.objfun),
                        }
                    # print("START DEBUG ... ")
                    # print(self.opti_solver.debug.g_describe(1))
                    # print(self.opti_solver.debug.x_describe(1))
                    # print('ang_tp2p23 :', self.opti_solver.debug.value(self.ang_tp2p23))
                    # print('ang_tp1p14 :', self.opti_solver.debug.value(self.ang_tp1p14))

                    solution_stats = self.opti_solver.stats()
                    if 'iterations' not in solution_stats: continue
                    if rich_output: print("Solution\n", solution)

                    obj_min = np.array(solution_stats['iterations']['obj']).min()
                    solution['obj_min'] = obj_min
                    solutions.append(solution)

                    if (solution['state'] == 'S') and (best_solution['state'] == 'F'): best_solution = solution
                    if (solution['objfun'] < best_solution['objfun']) and (best_solution['state'] == solution['state']): best_solution = solution

                    if solution['state'] == 'S': solution_succ_num += 1
                    if solution_succ_num > total_intervals/2: return solutions, best_solution
        return solutions, best_solution

    def build_grasp(self, solution):
        if solution['state'] == 'F': return [0, 0, 0, 0, 0], 'NNNNNNNNNN'

        # unpack result here
        ut1, vt1, th1 = solution['ut1'], solution['vt1'], solution['th1']
        ut2, vt2, th2 = solution['ut2'], solution['vt2'], solution['th2']
        ut3, vt3, th3 = solution['ut3'], solution['vt3'], solution['th3']
        ut4, vt4, th4 = solution['ut4'], solution['vt4'], solution['th4']
        dis_open = solution['dis_open']

        del_open_u1 = - dis_open * np.sin(th1)
        del_open_v1 = + dis_open * np.cos(th1)
        del_open_u2 = - dis_open * np.sin(th2)
        del_open_v2 = + dis_open * np.cos(th2)
        del_open_u3 = - dis_open * np.sin(th3)
        del_open_v3 = + dis_open * np.cos(th3)
        del_open_u4 = - dis_open * np.sin(th4)
        del_open_v4 = + dis_open * np.cos(th4)
        up1 = ut1 + del_open_u1
        vp1 = vt1 + del_open_v1
        up2 = ut2 + del_open_u2
        vp2 = vt2 + del_open_v2
        up3 = ut3 + del_open_u3
        vp3 = vt3 + del_open_v3
        up4 = ut4 + del_open_u4
        vp4 = vt4 + del_open_v4

        point_center = np.array([[(up1 + up2 + up3 + up4)/4, (vp1 + vp2 + vp3 + vp4)/4]])

        p_u, p_v = point_center[0, 0], point_center[0, 1]
        if self.gripper_type == "FishGri_4F": orientation = th1 if np.cos(th1) > 0 else np.arctan(np.sin(th1)/np.cos(th1))
        if self.gripper_type == "RochuGr_4F": orientation = th1 % (np.pi/2) - np.pi/4 # for convinience
        grsap_opening = dis_open
        palm_position = self.init_palm_dists[0]/grsap_opening
        grasp_params = [p_u, p_v, orientation, grsap_opening, palm_position]
        return grasp_params, self.gripper_type

    @staticmethod
    def plot_result(certainty_map, solution = None, rgb_image=None):
        """
        Convenient function for visualization
        """
        plt.figure("solution")
        if rgb_image is not None:
            plt.imshow(rgb_image)
            # certainty_map = np.ma.masked_where(certainty_map < 0.01, certainty_map)
            certainty_map = np.clip(certainty_map, 0, 1)
            plt.imshow(certainty_map, cmap='jet', alpha=certainty_map*0.8)
        else:
            plot = plt.matshow(certainty_map)
            plt.colorbar(plot)

        if solution['objfun'] > 1e8:
            plt.text(10, 20, 'No solution', fontdict=dict(fontsize=10, color='w', weight='bold'),
                bbox=dict(fill=True, color='k', linewidth=1, alpha=0.8))
            return plt

        # unpack result here
        ut1, vt1, th1 = solution['ut1'], solution['vt1'], solution['th1']
        ut2, vt2, th2 = solution['ut2'], solution['vt2'], solution['th2']
        ut3, vt3, th3 = solution['ut3'], solution['vt3'], solution['th3']
        ut4, vt4, th4 = solution['ut4'], solution['vt4'], solution['th4']
        dis_open = solution['dis_open']

        del_open_u1 = - dis_open * np.sin(th1)
        del_open_v1 = + dis_open * np.cos(th1)
        del_open_u2 = - dis_open * np.sin(th2)
        del_open_v2 = + dis_open * np.cos(th2)
        del_open_u3 = - dis_open * np.sin(th3)
        del_open_v3 = + dis_open * np.cos(th3)
        del_open_u4 = - dis_open * np.sin(th4)
        del_open_v4 = + dis_open * np.cos(th4)
        up1 = ut1 + del_open_u1
        vp1 = vt1 + del_open_v1
        up2 = ut2 + del_open_u2
        vp2 = vt2 + del_open_v2
        up3 = ut3 + del_open_u3
        vp3 = vt3 + del_open_v3
        up4 = ut4 + del_open_u4
        vp4 = vt4 + del_open_v4

        points_palms = np.array([[up1, vp1], [up2, vp2], [up3, vp3], [up4, vp4], [up1, vp1]])
        points_tp1 = np.array([[ut1, vt1], [up1, vp1]])
        points_tp2 = np.array([[ut2, vt2], [up2, vp2]])
        points_tp3 = np.array([[ut3, vt3], [up3, vp3]])
        points_tp4 = np.array([[ut4, vt4], [up4, vp4]])

        point_center = np.array([[(up1 + up2 + up3 + up4)/4, (vp1 + vp2 + vp3 + vp4)/4]])
        sta, cer, obj = solution['state'], solution['obj_cer'], solution['objfun']

        # start ploting
        plt.plot(points_palms[:, 1], points_palms[:, 0], '--', linewidth=2,  color='b', zorder=0)
        plt.plot(points_tp1[:, 1], points_tp1[:, 0], linewidth=4, color='k' if sta == 'S' else 'r', zorder=5)
        plt.plot(points_tp2[:, 1], points_tp2[:, 0], linewidth=4, color='k' if sta == 'S' else 'r', zorder=5)
        plt.plot(points_tp3[:, 1], points_tp3[:, 0], linewidth=4, color='k' if sta == 'S' else 'r', zorder=5)
        plt.plot(points_tp4[:, 1], points_tp4[:, 0], linewidth=4, color='k' if sta == 'S' else 'r', zorder=5)
        plt.scatter(point_center[:, 1], point_center[:, 0], s = 10, color='w', zorder=10)
        plt.text(10, 20, '{}: {:.1f},{:.1f}'.format(sta, cer, obj), fontdict=dict(fontsize=10, color='w', weight='bold'),
            bbox=dict(fill=True, color='k', linewidth=1, alpha=0.8))
        # plt.show()
        return plt

# ###
# ### Try it now !
# ###
# gripper_optimizer = GraspOptimSolver_4F(gripper_type='FishGrip4F') # FishGrip4F Rochu4F
# gripper_optimizer.refresh_mapsfuns(angle_map, certainty_map)
# gripper_optimizer.init_solver()

# sol_pack, sol_final, sol_best = gripper_optimizer.iterate_solver(\
#         max_iter=500, total_intervals=3, trials_each_interval=10, rich_output=False, decrease_rate=0.8)

# # for sol_pack_i in sol_pack: gripper_optimizer.plot_result(angle_map, sol_pack_i)

# best_objs = []
# for sol_pack_i in sol_pack:
#     best_objs.append(sol_pack_i['objfun'])
#     best_objs.append(sol_pack_i['obj_min'])
#     best_objs.append(sol_pack_i['state'])
#     print('\n\nsol_pack:', sol_pack_i)

# if(len(sol_pack) > 0): gripper_optimizer.plot_result_4F(angle_map, sol_best)
# print('\n\nsol_best:', sol_best)
# if(len(sol_pack) > 0): gripper_optimizer.plot_result_4F(angle_map, sol_final)
# print('\n\nsol_final:', sol_final)


# Imports

from __future__ import print_function
import sys
import cv2
from random import randint
import pandas as pd
import numpy as np
import os
import subprocess
from IPython.display import Video, display

# My working directory for the data, change it to fit your computer

working_dir = "/Users/camrogers/Downloads/Computer Vision/NFL Kaggle Data/"

# This is the video labels for each frame with the players and bounding boxes

vid_labels = pd.read_csv(working_dir + "train_labels.csv")

# Types of tracking algorithms that could be used

def create_tracker_csrt():
  tracker = cv2.legacy.TrackerCSRT_create()
  return tracker

# IOU score function

def bb_intersection_over_union(boxA, boxB):
    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    # compute the area of intersection rectangle
    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the interesection area
    iou = interArea / float(boxAArea + boxBArea - interArea)
    # return the intersection over union value
    return iou

# Function to train CSRT tracker on a video, compute IOU score, and display video with tracker's predictions

def get_iou_vid(videoPath, trackerType):
    # Create a video capture object to read videos

    cap = cv2.VideoCapture(working_dir+"/train/" + videoPath)

    # Read first frame

    success, frame = cap.read()

    # quit if unable to read the video file

    if not success:
      print('Failed to read video')
      sys.exit(1)

    ## List for boxes, colors of boxes, and player labels

    bboxes = []
    colors = []
    labels = []

    # First frame may be a black screen, this loop goes to the next frame until it's not black

    while True:
        if np.max(frame) == 0:
            success, frame = cap.read()
            if not success:
                print('Failed to read video')
                sys.exit(1)
        else:
            break

    # Bounding boxes of first frame of the given video

    bbox_df = vid_labels.query("video == @videoPath & frame == 1")

    for i in range(bbox_df.shape[0]):
        bboxes.append([bbox_df.iloc[i, 6], bbox_df.iloc[i, 8], bbox_df.iloc[i, 7], bbox_df.iloc[i, 9]])
        colors.append((randint(0, 255), randint(0, 255), randint(0, 255)))
        labels.append(bbox_df.iloc[i, 5])

    print('Selected bounding boxes {}'.format(bboxes))

    # Create MultiTracker object

    multiTracker = cv2.legacy_MultiTracker.create()

    # Initialize MultiTracker

    for bbox in bboxes:
      multiTracker.add(create_tracker_csrt(), frame, bbox)

    # j is a counter representing the current frame that the video is on

    j = 1

    # List of iou scores for each frame/player bounding box

    iou_scores = []

    # Updating the tracker

    while cap.isOpened():
      success, frame = cap.read()
      if not success:
          break
      # get updated location of objects in subsequent frames
      success, boxes = multiTracker.update(frame)
      j += 1
      # draw tracked objects
      for i, newbox in enumerate(boxes):
        # label of particular player helmet

        lab = labels[i]

        # predicted bounding boxes

        p1 = (int(newbox[0]), int(newbox[1]))
        p2 = (int(newbox[0] + newbox[2]), int(newbox[1] + newbox[3]))

        # drawing predictions on video

        #cv2.rectangle(frame, p1, p2, colors[i], 2, 1)
        #cv2.putText(frame, labels[i], (p1[0], max(0, p2[1] - 5)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color=colors[i], thickness=2)

        # getting true box labels

        bb_true = vid_labels.query("video == @videoPath & frame == @j & label == @lab")

        # player may drop out of video and bounding box no longer exists, just skip that box

        if bb_true.shape[0] == 0:
          continue

        # true bounding box from train_labels

        p1_true = (int(bb_true.iloc[0, 6]), int(bb_true.iloc[0, 8]))
        p2_true = (int(bb_true.iloc[0, 6]+bb_true.iloc[0, 7]), int(bb_true.iloc[0, 8]+bb_true.iloc[0, 9]))

        # IOU score for particular box/frame

        iou = bb_intersection_over_union([p1_true[0], p1_true[1], p2_true[0], p2_true[1]],
                                     [p1[0], p1[1], p2[0], p2[1]])
        iou_scores.append(iou)

      # show frame
      cv2.imshow('MultiTracker', frame)
      # quit on ESC button
      if cv2.waitKey(1) & 0xFF == 27:  # Esc pressed
          break

    cv2.destroyAllWindows()

    # Average IOU score for video

    print(np.mean(iou_scores))

    return np.mean(iou_scores)

# Taking random subset of 15 videos to test our tracker on (only used 15 because of runtime concerns)

video_names_end = vid_labels["video"].str.contains("Endzone")
video_names = vid_labels["video"][video_names_end].unique()

rand_idx = np.random.randint(0, len(video_names), 15)

video_names = video_names[rand_idx]

iou_scores = []

for name in video_names:
  iou = get_iou_vid(name, "CSRT")
  iou_scores.append(iou)
  print(iou)

np.mean(iou_scores) # About 0.42


import Chemins
from Utils import UTILS_Adaptations
from Utils.UTILS_Traduction import _
import wx
from Ctrl import CTRL_Bouton_image

from Ol import OL_Agrements
import GestionDB

class Panel(wx.Panel):
    def __init__(self, parent, IDactivite=None, nouvelleActivite=False):
        wx.Panel.__init__(self, parent, id=-1, name="panel_agrements", style=wx.TAB_TRAVERSAL)
        self.parent = parent
        self.IDactivite = IDactivite

        self.staticbox_agrements_staticbox = wx.StaticBox(self, -1, _(u"Agréments de l'activité"))
        self.radio_aucun = wx.RadioButton(self, -1, u"")
        self.radio_aucun.SetValue(True)
        self.label_aucun = wx.StaticText(self, -1, _(u"Aucun agrément"))
        self.radio_unique = wx.RadioButton(self, -1, u"")
        self.label_unique = wx.StaticText(self, -1, _(u"Agrément unique :"))
        self.ctrl_agrement_unique = wx.TextCtrl(self, -1, u"")
        self.radio_multiples = wx.RadioButton(self, -1, u"")
        self.label_multiples = wx.StaticText(self, -1, _(u"Agréments multiples :"))

        self.ctrl_agrements = OL_Agrements.ListView(self, IDactivite=self.IDactivite, id=-1, name="OL_agrements", style=wx.LC_REPORT|wx.SUNKEN_BORDER|wx.LC_SINGLE_SEL|wx.LC_HRULES|wx.LC_VRULES)
        self.ctrl_agrements.MAJ()

        self.bouton_ajouter = wx.BitmapButton(self, -1, wx.Bitmap(Chemins.GetStaticPath(u"Images/16x16/Ajouter.png"), wx.BITMAP_TYPE_ANY))
        self.bouton_modifier = wx.BitmapButton(self, -1, wx.Bitmap(Chemins.GetStaticPath(u"Images/16x16/Modifier.png"), wx.BITMAP_TYPE_ANY))
        self.bouton_supprimer = wx.BitmapButton(self, -1, wx.Bitmap(Chemins.GetStaticPath(u"Images/16x16/Supprimer.png"), wx.BITMAP_TYPE_ANY))

        self.__set_properties()
        self.__do_layout()

        self.Bind(wx.EVT_RADIOBUTTON, self.OnRadio, self.radio_aucun)
        self.Bind(wx.EVT_RADIOBUTTON, self.OnRadio, self.radio_unique)
        self.Bind(wx.EVT_RADIOBUTTON, self.OnRadio, self.radio_multiples)
        self.Bind(wx.EVT_BUTTON, self.OnBoutonAjouter, self.bouton_ajouter)
        self.Bind(wx.EVT_BUTTON, self.OnBoutonModifier, self.bouton_modifier)
        self.Bind(wx.EVT_BUTTON, self.OnBoutonSupprimer, self.bouton_supprimer)

        # Importation
        if self.IDactivite != None :
            self.Importation()

        self.OnRadio(None)

    def __set_properties(self):
        self.radio_unique.SetToolTip(wx.ToolTip(_(u"Sélectionnez 'Agrément unique' si l'activité possède un agrément unique")))
        self.ctrl_agrement_unique.SetMinSize((120, -1))
        self.ctrl_agrement_unique.SetToolTip(wx.ToolTip(_(u"Saisissez ici le numéro d'agrément unique")))
        self.radio_multiples.SetToolTip(wx.ToolTip(_(u"Sélectionnez 'Agréments multiples' si l'activité peut avoir plusieurs agréments")))
        self.bouton_ajouter.SetToolTip(wx.ToolTip(_(u"Cliquez ici pour ajouter un agrément")))
        self.bouton_modifier.SetToolTip(wx.ToolTip(_(u"Cliquez ici pour modifier l'agrément selectionné dans la liste")))
        self.bouton_supprimer.SetToolTip(wx.ToolTip(_(u"Cliquez ici pour supprimer l'agrément selectionnné dans la liste")))

    def __do_layout(self):
        grid_sizer_base = wx.FlexGridSizer(rows=1, cols=1, vgap=10, hgap=10)
        staticbox_agrements = wx.StaticBoxSizer(self.staticbox_agrements_staticbox, wx.VERTICAL)
        grid_sizer_agrements = wx.FlexGridSizer(rows=4, cols=2, vgap=5, hgap=5)
        grid_sizer_multiples = wx.FlexGridSizer(rows=1, cols=2, vgap=5, hgap=5)
        grid_sizer_boutons = wx.FlexGridSizer(rows=3, cols=1, vgap=5, hgap=5)
        grid_sizer_unique = wx.FlexGridSizer(rows=1, cols=2, vgap=5, hgap=5)
        grid_sizer_agrements.Add(self.radio_aucun, 0, 0, 0)
        grid_sizer_agrements.Add(self.label_aucun, 0, wx.ALIGN_CENTER_VERTICAL, 0)
        grid_sizer_agrements.Add(self.radio_unique, 0, wx.ALIGN_CENTER_VERTICAL, 0)
        grid_sizer_unique.Add(self.label_unique, 0, wx.ALIGN_CENTER_VERTICAL, 0)
        grid_sizer_unique.Add(self.ctrl_agrement_unique, 0, 0, 0)
        grid_sizer_agrements.Add(grid_sizer_unique, 1, wx.EXPAND, 0)
        grid_sizer_agrements.Add(self.radio_multiples, 0, wx.ALIGN_CENTER_VERTICAL, 0)
        grid_sizer_agrements.Add(self.label_multiples, 0, wx.ALIGN_CENTER_VERTICAL, 0)
        grid_sizer_agrements.Add((15, 15), 0, wx.EXPAND, 0)
        grid_sizer_multiples.Add(self.ctrl_agrements, 1, wx.EXPAND, 0)
        grid_sizer_boutons.Add(self.bouton_ajouter, 0, 0, 0)
        grid_sizer_boutons.Add(self.bouton_modifier, 0, 0, 0)
        grid_sizer_boutons.Add(self.bouton_supprimer, 0, 0, 0)
        grid_sizer_multiples.Add(grid_sizer_boutons, 1, wx.EXPAND, 0)
        grid_sizer_multiples.AddGrowableRow(0)
        grid_sizer_multiples.AddGrowableCol(0)
        grid_sizer_agrements.Add(grid_sizer_multiples, 1, wx.EXPAND, 0)
        grid_sizer_agrements.AddGrowableRow(3)
        grid_sizer_agrements.AddGrowableCol(1)
        staticbox_agrements.Add(grid_sizer_agrements, 1, wx.ALL|wx.EXPAND, 5)
        grid_sizer_base.Add(staticbox_agrements, 1, wx.ALL|wx.EXPAND, 10)
        self.SetSizer(grid_sizer_base)
        grid_sizer_base.Fit(self)
        grid_sizer_base.AddGrowableRow(0)
        grid_sizer_base.AddGrowableCol(0)

    def OnRadio(self, event):
        if self.radio_aucun.GetValue() == 1 :
            self.ctrl_agrement_unique.Enable(False)
            self.ctrl_agrements.Enable(False)
            self.bouton_ajouter.Enable(False)
            self.bouton_modifier.Enable(False)
            self.bouton_supprimer.Enable(False)
        if self.radio_unique.GetValue() == 1 :
            self.ctrl_agrement_unique.Enable(True)
            self.ctrl_agrements.Enable(False)
            self.bouton_ajouter.Enable(False)
            self.bouton_modifier.Enable(False)
            self.bouton_supprimer.Enable(False)
            self.ctrl_agrement_unique.SetFocus()
        if self.radio_multiples.GetValue() == 1 :
            self.ctrl_agrement_unique.Enable(False)
            self.ctrl_agrements.Enable(True)
            self.bouton_ajouter.Enable(True)
            self.bouton_modifier.Enable(True)
            self.bouton_supprimer.Enable(True)

    def OnBoutonAjouter(self, event):
        self.ctrl_agrements.Ajouter(None)

    def OnBoutonModifier(self, event):
        self.ctrl_agrements.Modifier(None)

    def OnBoutonSupprimer(self, event):
        self.ctrl_agrements.Supprimer(None)

    def Importation(self):
        """ Importation des données """
        db = GestionDB.DB()
        req = """SELECT IDagrement, agrement, date_debut, date_fin
        FROM agrements WHERE IDactivite=%d;""" % self.IDactivite
        db.ExecuterReq(req)
        listeDonnees = db.ResultatReq()
        db.Close()
        if len(listeDonnees) == 0 : return
        # Unique
        if len(listeDonnees) == 1 :
            activite = listeDonnees[0]
            if activite[2] == "1977-01-01" and activite[3] == "2999-01-01" :
                self.radio_unique.SetValue(True)
                self.ctrl_agrement_unique.SetValue(activite[1])
                return
        # Multiples
        if len(listeDonnees) >= 1 :
            self.radio_multiples.SetValue(True)

    def Validation(self):
        # Aucun
        if self.radio_aucun.GetValue() == True :
            DB = GestionDB.DB()
            req = """SELECT IDagrement, IDactivite, agrement, date_debut, date_fin
            FROM agrements
            WHERE IDactivite=%d; """ % self.IDactivite
            DB.ExecuterReq(req)
            listeDonnees = DB.ResultatReq()
            DB.Close()
            if len(listeDonnees) > 0 :
                dlg = wx.MessageDialog(self, _(u"Vous avez sélectionné 'Aucun agrément' pour cette activité. Souhaitez-vous vraiment supprimer le ou les agréments précédemment saisis pour cette activité ?"), _(u"Suppression d'agréments"), wx.YES_NO|wx.NO_DEFAULT|wx.CANCEL|wx.ICON_INFORMATION)
                if dlg.ShowModal() != wx.ID_YES :
                    dlg.Destroy()
                    return False
                else:
                    dlg.Destroy()
        # Unique
        if self.radio_unique.GetValue() == True :
            if self.ctrl_agrement_unique.GetValue() == "" :
                dlg = wx.MessageDialog(self, _(u"Vous avez sélectionné 'Agrément unique' sans saisir de numéro d'agrément !"), _(u"Erreur de saisie"), wx.OK | wx.ICON_EXCLAMATION)
                dlg.ShowModal()
                dlg.Destroy()
                self.ctrl_agrement_unique.SetFocus()
                return False
        # Multiples
        if self.radio_multiples.GetValue() == True :
            nbreAgrements = len(self.ctrl_agrements.donnees)
            if nbreAgrements == 0 :
                dlg = wx.MessageDialog(self, _(u"Vous avez sélectionné 'Agréments multiples' sans saisir aucun numéro d'agrément !"), _(u"Erreur de saisie"), wx.OK | wx.ICON_EXCLAMATION)
                dlg.ShowModal()
                dlg.Destroy()
                return False
        # Validation
        return True

    def Sauvegarde(self):
        # Aucun
        if self.radio_aucun.GetValue() == True :
            DB = GestionDB.DB()
            req = """SELECT IDagrement, IDactivite, agrement, date_debut, date_fin
            FROM agrements
            WHERE IDactivite=%d; """ % self.IDactivite
            DB.ExecuterReq(req)
            listeDonnees = DB.ResultatReq()
            if len(listeDonnees) > 0 :
                # Effacement des agréments existants
                DB.ReqDEL("agrements", "IDactivite", self.IDactivite)
            DB.Close()
        # Unique
        if self.radio_unique.GetValue() == True :
            # Effacement des numéros d'agréments multiples existants
            DB = GestionDB.DB()
            req = """SELECT IDagrement, IDactivite, agrement, date_debut, date_fin
            FROM agrements
            WHERE date_debut<>'1977-01-01' AND date_fin<>'2999-01-01' AND IDactivite=%d; """ % self.IDactivite
            DB.ExecuterReq(req)
            listeDonnees = DB.ResultatReq()
            if len(listeDonnees) > 0 :
                # Effacement des agréments uniques existants
                for IDagrement, IDactivite, agrement, date_debut, date_fin in listeDonnees :
                    DB.ReqDEL("agrements", "IDagrement", IDagrement)
            # Enregistrement de l'agrément unique
            req = """SELECT IDagrement, IDactivite, agrement, date_debut, date_fin
            FROM agrements
            WHERE date_debut='1977-01-01' AND date_fin='2999-01-01' AND IDactivite=%d; """ % self.IDactivite
            DB.ExecuterReq(req)
            listeDonnees = DB.ResultatReq()
            if len(listeDonnees) > 0 :
                IDagrement = listeDonnees[0][0]
            else:
                IDagrement = None
            listeDonneesSave = [
                ("IDactivite", self.IDactivite ),
                ("agrement", self.ctrl_agrement_unique.GetValue() ),
                ("date_debut", "1977-01-01"),
                ("date_fin", "2999-01-01"),
                ]
            if IDagrement == None :
                IDagrement = DB.ReqInsert("agrements", listeDonneesSave)
            else:
                DB.ReqMAJ("agrements", listeDonneesSave, "IDagrement", IDagrement)
            DB.Close()
        # Multiples
        if self.radio_multiples.GetValue() == True :
            DB = GestionDB.DB()
            req = """SELECT IDagrement, IDactivite, agrement, date_debut, date_fin
            FROM agrements
            WHERE date_debut='1977-01-01' AND date_fin='2999-01-01' AND IDactivite=%d; """ % self.IDactivite
            DB.ExecuterReq(req)
            listeDonnees = DB.ResultatReq()
            if len(listeDonnees) > 0 :
                # Effacement des agréments uniques existants
                for IDagrement, IDactivite, agrement, date_debut, date_fin in listeDonnees :
                    DB.ReqDEL("agrements", "IDagrement", IDagrement)
            DB.Close()

class MyFrame(wx.Frame):
    def __init__(self, *args, **kwds):
        wx.Frame.__init__(self, *args, **kwds)
        panel = wx.Panel(self, -1)
        sizer_1 = wx.BoxSizer(wx.VERTICAL)
        sizer_1.Add(panel, 1, wx.ALL|wx.EXPAND)
        self.SetSizer(sizer_1)
        self.ctrl= Panel(panel, IDactivite=1)
        sizer_2 = wx.BoxSizer(wx.VERTICAL)
        sizer_2.Add(self.ctrl, 1, wx.ALL|wx.EXPAND, 4)
        panel.SetSizer(sizer_2)
        self.Layout()
        self.CentreOnScreen()

if __name__ == '__main__':
    app = wx.App(0)
    #wx.InitAllImageHandlers()
    frame_1 = MyFrame(None, -1, _(u"TEST"), size=(700, 500))
    app.SetTopWindow(frame_1)
    frame_1.Show()
    app.MainLoop()


#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   pe_topN.py
@Time    :   2020/12/21 15:45:38
@Author  :   Jeffrey Wang
@Version :   1.0
@Contact :   shwangjj@163.com
@Desc    :   定期找出PE最小的10支股票

- 加载2019年所有的股票
- 每10天，找出pe最小的10个股票清单
- 输出这个股票清单
* 用name来检索datas

'''
import backtrader as bt
import os
import datetime
import pandas as pd

DATA_ROOT = "c:/fd_data/tushare/"

# 加载股票日线数据文件
def get_stock_daily_data(stock_id="600016.SH", start="20190101", end="20191231"):
    """从本地股票日线csv文件中取得DataFeed对象

    Args:
        stock_id (str, optional): tushare格式的股票代码. Defaults to "600016.SH".

    Raises:
        Exception: 本地文件无法找到

    Returns:
        [type]: backtrader.feeds.pandafeed.PandasData对象
    """
    # 日期格式转换
    dt_start = datetime.datetime.strptime(start, "%Y%m%d")
    dt_end = datetime.datetime.strptime(end, "%Y%m%d")

    # 文件存放地址
    file_path = DATA_ROOT + 'daily_bar/' + stock_id + ".csv"
    if not os.path.exists(file_path):
        print("数据源文件未找到！" + file_path)
        raise Exception("数据源文件未找到！" + file_path)

    # 将csv文件转为pandas.dataframe
    df = pd.read_csv(filepath_or_buffer=file_path)
    # 按日期先后排序
    df.sort_values(by=["trade_date"], ascending=True, inplace=True)
    # 将日期列，设置成index
    df.index = pd.to_datetime(df.trade_date, format='%Y%m%d')
    # 增加一列openinterest
    df['openinterest'] = 0.00
    # 取出特定的列
    df = df[['open', 'high', 'low', 'close', 'vol', 'openinterest']]
    # 列名修改成指定的
    df.rename(columns={"vol": "volume"}, inplace=True)

    data = bt.feeds.PandasData(dataname=df, fromdate=dt_start, todate=dt_end)
    return data

class BestPeStrategy(bt.Strategy):
    """
    最小PE策略

    """

    def log(self, txt, dt=None):
        """ log信息的功能 """
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # self.dataopen = self.datas[0].open
        self.dataopen = self.data.get("000001.SZ").open

    def next(self):
        # 按name取数据datafeed
        self.log(self.dataopen[0])

if __name__ == '__main__':
    # 初始化引擎
    cerebro = bt.Cerebro()

    # 给Cebro引擎添加策略
    cerebro.addstrategy(BestPeStrategy)

    # 设置初始资金：
    cerebro.broker.setcash(100000.0)  # 10万元

    # 从csv文件加载数据
    start = "20190101"
    end = "20190331"
    data = get_stock_daily_data(stock_id="000001.SZ", start=start, end=end)
    cerebro.adddata(data, name="000001.SZ")
    data = get_stock_daily_data(stock_id="000002.SZ", start=start, end=end)
    cerebro.adddata(data, name="000002.SZ")

    # 回测启动运行
    cerebro.run()


from flask import Flask, request, render_template, redirect
from db import Database

app = Flask(__name__)

@app.route('/', methods=['POST', 'GET'])
def main():
    db = Database()
    items = db.show_db()
    return render_template('index.html', items=items)

@app.route('/cari', methods=['POST'])
def cari():
    db = Database()
    items = db.cari(request.form['keyword'])
    return render_template('cari.html', items=items)

@app.route('/new', methods=['POST'])
def add():
    return render_template('new.html')

@app.route('/update', methods=['POST'])
def apply():
    db = Database()
    code = request.form['code']
    name = request.form['name']
    active = request.form['active']
    db.new(code, name, active)
    return redirect('/')

app.run(debug=True)


from django.urls import path
from . import views

urlpatterns = [
    path('helloworld/', views.helloWorld),
    path('', views.accountList, name='account-list'),

    path('categorias/', views.categoryList, name='category-list'),
    path('categoria/<int:id>', views.categoryView, name='category-view'),
    path('editarCategoria/<int:id>', views.editCategory, name='edit-category'),
    path('deletarCategoria/<int:id>', views.deleteCategory, name='delete-category'),
    path('novaCategoria/', views.newCategory, name='new-category'),

    path('lancamentos/', views.accountList, name='account-list'),
    path('lancamento/<int:id>', views.accountView, name='account-view'),
    path('editarLancamento/<int:id>', views.editAccount, name='edit-account'),
    path('deletarLancamento/<int:id>', views.deleteAccount, name='delete-account'),
    path('novoLancamento/', views.newAccount, name='new-account')
]


import os
import numpy as np
import h5py
import errno
import MDAnalysis as mda
from cvae.CVAE import CVAE
#from keras import backend as K
import tensorflow.keras.backend as K
#from sklearn.cluster import DBSCAN
from cuml import DBSCAN as DBSCAN
import cuml
import cupy as cp
import numpy as np

import adios2

from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *

def triu_to_full(cm0):
    num_res = int(np.ceil((len(cm0) * 2) ** 0.5))
    iu1 = np.triu_indices(num_res, 1)

    cm_full = np.zeros((num_res, num_res))
    cm_full[iu1] = cm0
    cm_full.T[iu1] = cm0
    np.fill_diagonal(cm_full, 1)
    return cm_full

def read_adios_file(adios_file):
    cms = []
    positions = []
    f = adios2.open(adios_file, mode='r')
    for fs in f:
        cm = fs.read('contact_map')
        position = fs.read('positions')
        cms.append(cm)
        positions.append(position)
    f.close()
    return cms, positions

'''
def read_h5py_file(h5_file):
    cm_h5 = h5py.File(h5_file, 'r', libver='latest', swmr=True)
    return cm_h5[u'contact_maps']

def adios_cm_to_cvae(adios_cm_data_lists):
    cm_all = np.hstack(adios_cm_data_lists)
    big_cms = []
    big_positions = []
    for fn in cm_all:
        small = read_adios_file(fn)
        big_cms += small[0]
        big_positions += small[1]
    return big_cms, big_positions
'''

def adios_to_cvae(fn, lastN=2000):
    with adios2.open(fn, "r") as fr:
        n = fr.steps()
        vars = fr.available_variables()
        # print("vars = ", vars)
        print("n = ", n)
        results = {}
        for v in ['contact_map', 'positions', 'md5', 'velocities']:
            print(v)
            if(v != 'md5'):
                shape = list(map(int, vars[v]['Shape'].split(",")))
                zs = list(np.zeros(len(shape), dtype='int'))
                results[v] = fr.read(v, zs, shape, 0, n)
                # print(results[v].shape)
            else:
                results[v] = fr.read_string(v,0,n)

    return results['contact_map'][-lastN:], results['positions'][-lastN:], results['md5'][-lastN:], results['velocities'][-lastN:]

def adios_to_cvae_var(fn, v, lastN=2000):
    with adios2.open(fn, "r") as fr:
        n = fr.steps()
        print("n = ", n)
        shape = list(map(int, vars[v]['Shape'].split(",")))
        zs = list(np.zeros(len(shape), dtype='int'))
        result = fr.read(v, zs, shape, 0, n)
    return result[-lastN:]

def adios_list_to_cvae(fl, lastN=2000):
    cm = []
    positions = []
    md5 = []
    velocities = []
    for fn in fl:
        t = adios_to_cvae(fn, lastN)
        cm.append(t[0])
        positions.append(t[1])
        md5.append(t[2])
        velocities.append(t[3])
    return np.concatenate(cm), np.concatenate(positions), np.concatenate(md5), np.concatenate(velocities)

def adios_list_to_cvae_var(fl, v, lastN=2000):
    results = []
    for fn in fl:
        t = adios_to_cvae_var(fn, v, lastN)
        results.append(t)
    return np.concatenate(results)

'''

def adios_to_cvae(adios_file_lists):
    big_cm = []
    big_positions = []
    for fn in adios_file_lists:
        small = read_adios_file(fn)
        cms = small[0]
        cms2 = [triu_to_full(cm_data) for cm_data in cms]
        big_cm += cms2
        big_positions += small[1]

    cm_data_full = np.array(big_cm)

    pad_f = lambda x: (0,0) if x%2 == 0 else (0,1)
    padding_buffer = [(0,0)]
    for x in cm_data_full.shape[1:]:
        padding_buffer.append(pad_f(x))
    cm_data_full = np.pad(cm_data_full, padding_buffer, mode='constant')

    # reshape matrix to 4d tensor
    cvae_input = cm_data_full.reshape(cm_data_full.shape + (1,))

    return cvae_input, big_positions
'''

'''
def cm_to_cvae(cm_data_lists):
    """
    A function converting the 2d upper triangle information of contact maps
    read from hdf5 file to full contact map and reshape to the format ready
    for cvae
    """
    cm_all = np.hstack(cm_data_lists)

    # transfer upper triangle to full matrix
    cm_data_full = np.array([triu_to_full(cm_data) for cm_data in cm_all.T])

    # padding if odd dimension occurs in image
    pad_f = lambda x: (0,0) if x%2 == 0 else (0,1)
    padding_buffer = [(0,0)]
    for x in cm_data_full.shape[1:]:
        padding_buffer.append(pad_f(x))
    cm_data_full = np.pad(cm_data_full, padding_buffer, mode='constant')

    # reshape matrix to 4d tensor
    cvae_input = cm_data_full.reshape(cm_data_full.shape + (1,))

    return cvae_input
'''

def stamp_to_time(stamp):
    import datetime
    return datetime.datetime.fromtimestamp(stamp).strftime('%Y-%m-%d %H:%M:%S')

'''
def find_frame(traj_dict, frame_number=0):
    local_frame = frame_number
    for key in sorted(traj_dict.keys()):
        if local_frame - int(traj_dict[key]) < 0:
            dir_name = os.path.dirname(key)
            traj_file = os.path.join(dir_name, 'output.dcd')
            return traj_file, local_frame
        else:
            local_frame -= int(traj_dict[key])
    raise Exception('frame %d should not exceed the total number of frames, %d' % (frame_number, sum(np.array(traj_dict.values()).astype(int))))

def write_pdb_frame(traj_file, pdb_file, frame_number, output_pdb):
    mda_traj = mda.Universe(pdb_file, traj_file)
    mda_traj.trajectory[frame_number]
    PDB = mda.Writer(output_pdb)
    PDB.write(mda_traj.atoms)
    return output_pdb
'''

def write_pdb_frame(frame, original_pdb, output_pdb_fn):
    pdb = PDBFile(original_pdb)
    f = open(output_pdb_fn, 'w')
    PDBFile.writeFile(pdb.getTopology(), frame, f)
    f.close()

def make_dir_p(path_name):
    try:
        os.mkdir(path_name)
    except OSError as exc:
        if exc.errno != errno.EEXIST:
            raise
        pass

def outliers_from_cvae(model_weight, cvae_input, hyper_dim=3, eps=0.35):
    #os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
    #os.environ["CUDA_VISIBLE_DEVICES"]=str(0)
    cvae = CVAE(cvae_input.shape[1:], hyper_dim)
    cvae.model.load_weights(model_weight)
    cm_predict = cvae.return_embeddings(cvae_input)
    print("In outliers_from_cvae")
    print("cm_predict")
    print(cm_predict)
    db = DBSCAN(eps=eps, min_samples=10).fit(cm_predict)
    db_label = db.labels_
    print("db_label")
    print(list(db_label))
    outlier_list = np.where(db_label == -1)
    K.clear_session()
    return outlier_list

def predict_from_cvae(model_weight, cvae_input, hyper_dim=3):
    #os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
    #os.environ["CUDA_VISIBLE_DEVICES"]=str(0)
    print("predict 1:", type(cvae_input))
    print("predict 1:", cvae_input.shape)
    cvae = CVAE(cvae_input.shape[1:], hyper_dim)
    cvae.model.load_weights(model_weight)
    cm_predict = cvae.return_embeddings(cvae_input)
    del cvae
    K.clear_session()
    return cm_predict

def outliers_from_latent(cm_predict, eps=0.35):
    #print(f"In outliers_from_latent eps={eps}")
    #print("cm_predict")
    #print(type(cm_predict))
    #print(cm_predict)
    cm_predict = cp.asarray(cm_predict)
    #print(type(cm_predict))
    #print(cm_predict)
    db = DBSCAN(eps=eps, min_samples=10, max_mbytes_per_batch=500).fit(cm_predict)
    #print(dir(db))
    db_label = db.labels_.to_array()
    #print("db_label")
    #print(list(db_label))
    outlier_list = np.where(db_label == -1)
    return outlier_list


#!/usr/bin/env python3
from multiprocessing.pool import ThreadPool as Pool
import matplotlib.pyplot as plt
import subprocess as sp
import platform
import time
import os

class address_checker():
    """This is a class to evaluate responsativity of multiple addresses to a
       ping and find addresses with same host but different networks addresses
       that have different responses to a ping.

       This library has only been tested on Windows OS.

    :param network_addresses: a list of network addresses to check such as
                              ['192.168.1.', '192.168.2.']. There should at
                              least be two network addresses.
    :type network_addresses: list

    :param host_addresses: a list of integers containing all host addresses,
                           defaults to [0, ..., 255]
    :type host_addresses: list, optional

    :param host_unwanted: a list of host addresses that should not be evaluates
                          , defaults to []
    :type host_unwanted: list, optional

    :param number_threads: number of threads to accelerate the processes of
                           pinging all addresses. For optimal performance,
                           utilize self.optimal_thread_number() for optimal
                           thread number, defaults to 130
    :type number_threads: int, optional

    :param n_echos: specifies the number of echo request messages sent,
                    defaults to 1
    :type n_echos: int, optional

    :param wait: specifies the amount of time, in milliseconds, to wait for
                 the echo, defaults to 2 ms
    :type wait: int, optional

    :param n_attempts: number of attempts in case an address is not responsive
                        to the ping, defaults to 2
    :type n_attempts: int, optional

    Examples
    --------
    >>> import address_checker from networking
    >>> host_unwanted = [15, 56]
    >>> network_addresses = ['192.168.1.', '192.168.2.']
    >>> c = address_checker(network_addresses)
    >>> c.ping_all()

    """

    def __init__(self, network_addresses, host_addresses=list(range(256)),
                 host_unwanted=[], number_threads=130, n_echos=1, wait=2,
                 n_attempts=2):
        """Constructor method.
        """
        # Assign properties
        self.network_addresses = network_addresses
        self.host_addresses = host_addresses
        self.host_unwanted = host_unwanted
        self.number_threads = number_threads
        self.n_attempts = n_attempts
        self.n_echos = n_echos
        self.wait = wait

        # Option for the number of packets as a function of OS
        if platform.system().lower() == 'windows':
            self.param = '-n'
        else:
            self.param = '-c'

        # Check if inputs are properly formatted
        self._check_inputs()

    def _check_inputs(self):
        '''Evaluates if inputs are the correct format. This is evaluated when
           the object is created. The attribute is also invoked before any
           major operation (ping multiple addresses) just in case the user
           changed the variables.
        '''

        # Check if OS is Windows. It should run for Linux, but was not tested
        if platform.system().lower() != 'windows':
            raise TypeError(
                'This library has not been verified for non-Windows OS')

        # Check if all list inputs are properly formatted
        for name in ['network_addresses', 'host_addresses', 'host_unwanted']:
            values = getattr(self, name)

            # Check if inputs are lists
            if not isinstance((values), (list, tuple)):
                raise TypeError(
                    name + " should be an indexable format such as a list")

            # Check for empty lists
            if name in ['network_addresses', 'host_addresses']:
                if len((values)) == 0:
                    raise TypeError(name + " should not be empty")

            if name == 'network_addresses':
                # At least two network addresses should be provided
                if len(values) < 2:
                    raise TypeError(
                        name + " should have at least two network addresses.")

                for value in values:
                    # network_addresses should consist only of strings
                    if not isinstance(value, (str)):
                        raise TypeError(name + " should consists of strings" +
                                               " such as '192.168.1.'")

                    split_value = value.split('.')
                    # network_addresses should follow a quad-dotted address
                    if len(split_value) != 4:
                        raise TypeError(name + " should consists of strings" +
                                               " such as '192.168.1.'")

                    # the host address in network_addresses should be missing
                    if split_value[-1] != '':
                        raise TypeError(
                            name + " should be missing the last octet of IP")

                    # Each octet should be less than 256
                    if not all(int(oct) < 256 for oct in split_value[:-1]):
                        raise TypeError("Each octet should be lower than 256")

            # host_addresses and host_unwanted should consists only of integers
            else:
                if not all(isinstance(value, (int)) for value in values):
                    if not all((val >= 0 and val < 256) for val in values):
                        raise TypeError(name + "should consists of positive" +
                                               " integers lower than 256")

        # Check if all list inputs are properly formmated
        for name in ['number_threads', 'n_echos', 'n_attempts']:
            values = getattr(self, name)

            if not isinstance((values), (int)):
                raise TypeError(name + " should be an integer")

            if values < 1:
                raise TypeError(
                    name + " should be an integer equal to or more than one")

    def show(self):
        """Prints all attributes from object.
        """
        attrs = vars(self)
        print(''.join("%s: %s\n" % item for item in attrs.items()))

    def ping(self, address):
        """Pings address.

        :param address: a string in dot-decimal notation is expected
        :type address: string

        :return: True if an echo is received
        :rtype: bool

        Notes:
        https://docs.microsoft.com/en-us/windows-server/administration/windows-commands/ping
        """
        self._check_inputs()

        exit_code = False
        attempts = 0

        while attempts < self.n_attempts and not exit_code:
            # Building the command. Ex: "ping -n 2 -w 2 [IP]"
            # (sp.call returns False positives)
            ping = sp.Popen("ping {} {} -w {} {}".format(self.param,
                            self.n_echos, self.wait, address),
                            stdout=sp.PIPE, stderr=sp.PIPE)

            # If adress responds, exit_code is True
            exit_code = ping.wait() == 0
            attempts += 1

        return exit_code

    def _ping_networks(self, host_address):
        """Pings the same host address for all specified network addresses.

        :param host_address: the host address (last octect of dot-decimal
                             notation). It should be positive and less than 256
        :type host_address: int

        :return: True if an echo is received
        :rtype: bool

        """
        if host_address not in self.host_unwanted:
            p = []
            for network_address in self.network_addresses:
                address = network_address + '%i' % host_address
                pi = self.ping(address)
                p.append(pi)

            if not all(pi == p[0] for pi in p):
                return host_address

    def ping_all(self, runtime=False):
        '''pings all addresses for all networks.

        :param runtime: if True, will runtime (default is False)
        :type runtime: bool, optional
        '''

        self._check_inputs()

        # Keep track of time
        if runtime:
            start_time = time.time()

        # Ping utilizing multithreading
        p = Pool(self.number_threads)
        unmatched_hosts = p.map(self._ping_networks, self.host_addresses)
        p.close()
        p.join()

        # Networks addresses with same ping results return None and are removed
        self.unmatched_hosts = [i for i in unmatched_hosts if i]

        # Find out time run
        if runtime:
            self.runtime = time.time() - start_time

    def optimal_thread_number(self, thread_range=list(range(10, 210, 10)),
                              plot=False):
        """Determines number of thread for best performance.

        :param address: a string in dot-decimal notation is expected
        :type address: string

        :return: (optimal runtime, optimal thread number)
        :rtype: float, int

        """
        runtimes = []
        for self.number_threads in thread_range:
            self.ping_all(runtime=True)
            runtimes.append(self.runtime)

        if plot:
            plt.figure()
            plt.plot(thread_range, runtimes)
            plt.xlabel('Number of threads')
            plt.ylabel('Runtime (s)')
            plt.show()

        optimal_runtime = min(runtimes)
        return optimal_runtime, (thread_range[runtimes.index(optimal_runtime)])


from typing import Dict, List

import numpy as np
import pandas as pd

import chimerapy.engine as cpe
from chimerapy.orchestrator import sink_node

from .data import YOLOFrame

def to_dataframe(results, frame_cnt, normalize=False):
    """
    Helper function to save results as pandas df
    Code adapted from method tojson() from YOLOv8 repo results.py
    """
    lst = []
    data = results.boxes.data.cpu().tolist()
    h, w = results.orig_shape if normalize else (1, 1)
    for i, row in enumerate(data):
        box = np.array([row[0] / w, row[1] / h, row[2] / w, row[3] / h])
        conf = row[4]
        id = int(row[5])
        name = results.names[id]
        result = {
            "frame_count": frame_cnt,
            "name": name,
            "class": id,
            "confidence": conf,
            "box": box,
        }
        if results.masks:
            xy = results.masks.xy[i]
            result["segments"] = np.array(
                [
                    (xy[:, 0] / w).tolist(),
                    (xy[:, 1] / h).tolist(),
                ]
            )
        if results.keypoints is not None:
            x, y, visible = (
                results.keypoints[i].data[0].cpu().unbind(dim=1)
            )  # torch Tensor
            result["keypoints"] = np.array(
                [
                    (x / w).tolist(),
                    (y / h).tolist(),
                    visible.tolist(),
                ]
            )
        lst.append(result)

    return pd.DataFrame(lst)

@sink_node(name="CPPipelines_YoloMultiSaveNode")
class MultiSaveNode(cpe.Node):
    """A node to save results from Yolov8 models

    Parameters
    ----------
    name : str, optional (default: 'MultiSaveNode')
        The name of the node
    source_key: str
        The name of the source video that needs its results saved
    frames_key: str, optional (default: 'frame')
        The key to access the frames in the video.
    filename: str, optional (default: 'yolo_results')
        The name of the file that results will be saved to
    file_format: str, optional (default: 'df')
        The format that results will be saved as. Two option available
        video (mp4, param: vid) and table (csv, param: df)
    fps: int, optional (default: 30)
        Video fps that can be manually set
    """

    def __init__(
        self,
        source_key: str,
        frames_key: str = "frame",
        name: str = "SaveNode",
        filename: str = "yolo_results",
        file_format: str = "df",
        fps: int = 30,
    ) -> None:
        self.source_key = source_key
        self.frames_key = frames_key
        self.format = file_format
        self.filename = filename
        self.fps = fps
        super().__init__(name=name)

    def step(self, data_chunks: Dict[str, cpe.DataChunk]) -> None:
        for _, data_chunk in data_chunks.items():
            frames: List[YOLOFrame] = data_chunk.get(self.frames_key)["value"]
            for frame in frames:
                if self.source_key == frame.src_id:
                    if "df" == self.format:
                        results = frame.result
                        frame_cnt = frame.frame_count
                        if results:
                            dfs = [
                                to_dataframe(result, frame_cnt)
                                for result in results
                            ]
                            df = pd.concat(dfs, ignore_index=True)
                            self.save_tabular(
                                self.filename + "-" + frame.src_id, df
                            )

                    elif "vid" == self.format:
                        img = frame.arr
                        if img.size > 0:
                            self.save_video(
                                self.filename + "-" + frame.src_id,
                                img,
                                self.fps,
                            )


import numpy as np
from mpi4py import MPI
#from scipy.interpolate import griddata, RegularGridInterpolator, NearestNDInterpolator, LinearNDInterpolator
from pymesher.models_1D import model
from pymesher.skeleton import Skeleton
from scipy.ndimage.filters import gaussian_filter
from math import pi, floor

# SET STUFF HERE.
scale = 6371000.
blockdir = '../MODEL_7_subJapan/'
DOMINANT_PERIOD = 15.0
ELEMENTS_PER_WAVELENGTH = 2.0
MIN_RADIUS =  5800000. # continue downward about 30 % #5771.0000 / 6371.000
REFERENCE_RADIUS = 6371.000
BLOCK_X = blockdir+"block_x"
BLOCK_Y = blockdir+"block_y"
BLOCK_Z = blockdir+"block_z"
output_model_name = '../Simute_2016_Japan_15s_flat_iso.e'
disc_style = 'keep_1D' #keep_1D, smooth
apply_gauss_smooth = True#False
gauss_stds = [2,2,4]

print('*'*60)
print('Interpolation script for a ses3d model to a cartesian salvus mesh.')
print('It is necessary to run this on as many cores as VARIABLES in the below list.')
print('*'*60)

# First item is name of variable in exodus file - second the filename
VARIABLES = [("VS",  blockdir+"dvs_ql6_1Hz_voigt_on_ref"),
             ("VP",  blockdir+"dvp_ql6_1Hz"),
             ("RHO", blockdir+"drho_ORIG"),
             ("QMU",blockdir+"ql_6")]

def smooth_ses3_blocks(var,nx,ny,nz,sig_x,sig_y,sig_z):
    # Given a 1-D array of elastic model parameters,
    # reshape by ses3d convention and smooth w Gaussian kernel

    var = var.reshape(nx,ny,nz)
    var = gaussian_filter(var,[sig_x,sig_y,sig_z])
    return var.ravel()

def get_cube(coord,x,y,z):
    # Given a coordinate tuple and a regular grid defined by x,y,z, return the
    # vertices or the cube containing the coordinate point.
    Dx = np.diff(x).max()
    Dy = np.diff(y).max()
    Dz = np.diff(z).max()

    indx_x = int(floor(((coord[0]-x.min())/Dx)))
    indx_y = int(floor(((coord[1]-y.min())/Dy)))
    indx_z = int(floor(((coord[2]-z.min())/Dz)))

    if indx_x < 0:
        indx_x = 0
        incr_x = 0
    elif indx_x >= len(x)-1:
        indx_x = len(x)-1
        incr_x = 0
    else:
        incr_x = 1

    if indx_y < 0:
        indx_y = 0
        incr_y = 0
    elif indx_y >= len(y)-1:
        indx_y = len(y)-1
        incr_y = 0
    else:
        incr_y = 1

    if indx_z < 0:
        indx_z = 0
        incr_z = 0
    elif indx_z >= len(z)-1:
        indx_z = len(z)-1
        incr_z = 0
    else:
        incr_z = 1
    # order:
    # point
    # 000
    # 100
    # 001
    # 101
    # 010
    # 110
    # 011
    # 111

    cube = [[indx_x, indx_y, indx_z],
            [indx_x+incr_x, indx_y, indx_z],
            [indx_x, indx_y, indx_z+incr_z],
            [indx_x+incr_x, indx_y, indx_z+incr_z],
            [indx_x, indx_y+incr_y, indx_z],
            [indx_x+incr_x, indx_y+incr_y, indx_z],
            [indx_x, indx_y+incr_y, indx_z+incr_z],
            [indx_x+incr_x, indx_y+incr_y, indx_z+incr_z]]

    return cube

def tril_int(point,points,values):

    # point is expected to be a tuple or array
    # (x,y,z), the new point
    # points is expected to be an array 8 x 3 (cube vertices)
    # values is expected to be an array of 8 values
    # If xyz are the coordinates of the vertices, expect the
    # points and values in the following order:

    # point   value
    # 000       0
    # 100       1
    # 001       2
    # 101       3
    # 010       4
    # 110       5
    # 011       6
    # 111       7

    x0 = points[0,0]
    x1 = points[1,0]

    y0 = points[0,1]
    y1 = points[4,1]

    z0 = points[0,2]
    z1 = points[2,2]

    # we have margins
    if x0 == x1:
        xd = 0.
    else:
        xd = (point[0]-x0) / (x1-x0)
    if y0 == y1:
        yd = 0.
    else:
        yd = (point[1]-y0) / (y1-y0)
    if z0 == z1:
        zd = 0.
    else:
        zd = (point[2]-z0) / (z1-z0)
    # print xd, yd, zd
    # x dimension

    c00 = values[0] * (1-xd) + values[1] * xd
    c01 = values[3] * (1-xd) + values[2] * xd
    c10 = values[4] * (1-xd) + values[5] * xd
    c11 = values[7] * (1-xd) + values[6] * xd

    # print c00,c01,c10,c11
    # y dimension
    c0 = c00 * (1-yd) + c10 * yd
    c1 = c01 * (1-yd) + c11 * yd
    # print c0,c1
    # z dimension
    c = c0 * (1-zd) + c1 * zd

    return c

# Begin of interpolation routine
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
if len(VARIABLES) != size:
    raise NotImplementedError("you have to run on as many cores as there are variables.")

# Read stuff.
theta = np.loadtxt(blockdir+"block_x", dtype=np.float64)
phi = np.loadtxt(blockdir+"block_y", dtype=np.float64)
r = np.loadtxt(blockdir+"block_z", dtype=np.float64)

nblocks = theta[0]
if phi[0] != nblocks or r[0] != nblocks:
    msg = 'Block files have different nr. of discontinuities.'
    raise ValueError(msg)
# ses3D has radius in km
r *= 1000.

if nblocks == 1:
    theta = theta[2:]
    phi = phi[2:]
    r = r[2:]
else:
    raise NotImplementedError(msg)

# Patch r to go to the surface which is easier for the mesher and which also
# will be the case for any real model.
# (only for tiny test model)
# r += (6371 - r.max())

# Normalize radius
# r /= REFERENCE_RADIUS

# initialize model object
mod = model.built_in("prem_ani")
print("Set up model object with 1-D background.")

# Get hmax
hmax = mod.get_edgelengths(
    dominant_period=DOMINANT_PERIOD,
    elements_per_wavelength=ELEMENTS_PER_WAVELENGTH)

hmax *= scale
discontinuities = mod.discontinuities*scale

# Only a chunk
full_sphere = False

# adapt discontinuities and hmax for min_radius
idx = discontinuities > MIN_RADIUS
ndisc = idx.sum() + 2

discontinuities_new = np.zeros(ndisc)
discontinuities_new[1] = MIN_RADIUS
discontinuities_new[-ndisc+2:] = discontinuities[idx]
#discontinuities = discontinuities_new[1:]

hmax_new = np.ones(ndisc-1)
hmax_new[-ndisc+2:] = hmax[-ndisc+2:]
#hmax = hmax_new[1:]

full_sphere = False

discontinuities = discontinuities_new[:]
hmax = hmax_new[:]

# The 1-D model may have discontinuities in the selected depth range

#==============================================
# Not sure I understand why we need this:
#==============================================
#discontinuities = discontinuities_new[:] # i.e. discontinuities[0] = 0
#hmax = hmax_new[:]    # i.e. hmax[0] = 1.0

# This doesn't work at the moment:
# elif disc_style == 'smooth': # No discontinuities; For a ses3D model that has
# # only one block region. Unfortunately, safest is then to use the smallest hmax
#      if nblocks > 1:
#         raise ValueError('Model contains discontinuities, cant use smooth')

#      discontinuities = np.array([MIN_RADIUS / REFERENCE_RADIUS, r[-1]])
#      hmax = np.array([1.,min(hmax)])
#      mod.nregions = 1

# Get the minimum span of phi. As it is not at the equator it is smaller.
#
# A --- B
# |   /
# | /
# C
# small circle radius = AB
#
# Calculate it at the smallest colatitude.

if disc_style == 'smooth':
    refinement = 'doubling_single_layer'
else:
    refinement = 'tripling'

_AB = np.sin(np.deg2rad(theta.min()))
# We'll build a spherical chunk at the equator and later rotate it.
sk = Skeleton.create_spherical_mesh(
    discontinuities, hmax, ndim=3,
    hmax_refinement=1.5,
    max_colat=[_AB * phi.ptp() * 0.6, theta.ptp() * 0.5],
    min_colat=[-_AB * phi.ptp() * 0.6, -theta.ptp() * 0.5],
    axisem=False,
    full_sphere=False,
    refinement_style=refinement,
    refinement_top_down=True)

m = sk.get_unstructured_mesh()

print('Created sperical mesh.')

#==============================================
# Does anything happen if we use this:
#==============================================
# find outer boundaries
side_set_mode = "spherical"
full_str = {True: '_full', False: '_chunk_z'}
side_set_mode += full_str[full_sphere]
m.find_side_sets(side_set_mode)

# Rotate to where the mesh should actually be located.
m.rotate_coordinates((0,theta.mean(), phi.mean()))
print('Rotated coordinates.')

theta = np.deg2rad(theta)
phi = np.deg2rad(phi)

# Putting coordinates inside the blocks
# 1   2   3   4   5   6   7
# 1   2.5 3.5 4.5 5.5 7

# Wait, wait...how about
# 0     1     2   3   4   5
# 0.5   1.5   2.5 3.5 4.5
#=================================
# I replaced this:
# ================================
# theta[1:-1] += np.diff(theta[1:])
# phi[1:-1] += np.diff(phi[1:])
# r[1:-1] += np.diff(r[1:])

# theta[-2] = theta[-1]
# phi[-2] = phi[-1]
# r[-2] = r[-1]

# theta = theta[:-1]
# phi = phi[:-1]
# r = r[:-1]

# d_theta = np.diff(theta[1:]).max()
# d_phi = np.diff(phi[1:]).max()
# d_r = np.diff(r[1:]).max()

# theta[0] -= 0.5 * d_theta
# phi[0] -= 0.5 * d_theta
# r[0] -= 0.5 * d_theta
# theta[-1] += 0.5 * d_theta
# phi[-1] += 0.5 * d_theta
# r[-1] += 0.5 * d_theta
# ================================
# With this:
# ================================

d_theta = np.diff(theta[1:]).max()
d_phi = np.diff(phi[1:]).max()
d_r = np.diff(r[1:]).max()

theta += 0.5 * d_theta
phi += 0.5 * d_phi
r += 0.5 * d_r

theta = theta[:-1]
phi = phi[:-1]
r = r[:-1]

thetav, phiv, rv = np.meshgrid(theta, phi, r)

#==============================================
# CARTESIAN coordinates here:
#==============================================
# # transform the blocks:
# x = rv * np.sin(thetav) * np.cos(phiv)
# y = rv * np.sin(thetav) * np.sin(phiv)
# z = rv * np.cos(thetav)

# transform mesh points to geographical coordinates,
# to locate them in ses3d grid
x_m = m.points[:, 0]
y_m = m.points[:, 1]
z_m = m.points[:, 2]

r_mesh = np.sqrt(np.power(x_m,2) +
                np.power(y_m,2) +
                np.power(z_m,2))
#==============================================
# Hmmmm...why is the arccos needed here. should
# be arcsin as we want to get latitude, not colat
# (but if I use arcsin, the resulting angles are
# somewhere in a completely weird place)
# Ahaaaaaa, ses3d uses colatitude :D :D :D nvm
#==============================================
theta_mesh = np.arccos(z_m / r_mesh)
phi_mesh = np.arctan2(y_m,x_m)

print('Transformed all mesh coordinates...')

npoints = len(m.points[:,0])
nx = len(theta)
ny = len(phi)
nz = len(r)

print('='*60)
print("Block grid extent:")
print(np.max(theta), np.max(phi), np.max(r))
print(np.min(theta), np.min(phi), np.min(r))
print("Salvus grid extent:")
print(np.max(theta_mesh), np.max(phi_mesh), np.max(r_mesh))
print(np.min(theta_mesh), np.min(phi_mesh), np.min(r_mesh))
print('='*60)
print("Salvus grid extent Cartesian:")
print(np.max(x_m), np.max(y_m), np.max(z_m))
print(np.min(x_m), np.min(y_m), np.min(z_m))
print('='*60)

(var, filename)= VARIABLES[rank]
print(var)
dat = np.loadtxt(filename, skiprows=2,dtype=np.float64)
if var in ['VP','VS','VSV','VSH']:
    print('Converting ses3d model (km/s) to (m/s)')
    dat *= 1000.

#==============================================
# Smoothing
#==============================================
if apply_gauss_smooth:

    dat = smooth_ses3_blocks(dat,len(theta),len(phi),len(r),
        gauss_stds[0],gauss_stds[1],gauss_stds[2])
    print("Applied Gaussian smoothing to " + var)

# walk through the mesh grid points
newvar = np.zeros(npoints)
for i in np.arange(npoints):
    if i%1e6 == 0:
        print("Finished %g points of %g." %(i,npoints))
# find cube and values from spherical regular grid
    point = [theta_mesh[i],phi_mesh[i],r_mesh[i]]
    inds = get_cube(point,theta,phi,r)

#==============================================
# CARTESIAN coordinates here:
#==============================================
    # point = [m.points[i,:]]
    #coords = np.array([(x[ix[0],ix[1],ix[2]],
    #          y[ix[0],ix[1],ix[2]],
    #          z[ix[0],ix[1],ix[2]]) for ix in inds])
#==============================================
    coords = np.array([(theta[ix[0]],
              phi[ix[1]],
              r[ix[2]]) for ix in inds])

    inds = np.array(inds)
    values = dat[inds[:,2]+inds[:,1]*nz+inds[:,0]*nz*ny]
    # print point[0]*180./pi,point[1]*180./pi,point[2]*REFERENCE_RADIUS
    # print inds
    # print coords
    # print values

# interpolate trilinearly on the spherical grid...cos the
# cartesian one is not regular

    newvar[i] = tril_int(point,coords,np.array(values))

comm.barrier()

if rank == 0:
    m.attach_field(var,newvar)

for i in range(1,size):

    if rank == i:
        comm.send(newvar,dest=0,tag=i)
    elif rank == 0:
        newvar = comm.recv(source=i,tag=i)
        print('Received data from rank %g' %i)

        varname = VARIABLES[i][0]

        m.attach_field(varname,newvar)

if rank == 0:
    m.attach_field('fluid', np.zeros(m.nelem))
    m.write_exodus(output_model_name)


from functools import wraps

def retry(max_retries: int):
    """Decorator to retry a flaky test a fix number of times before failing.

    Note: This does not implement any form of backoff so should not be use for tests which are flaky due to network
    issues.
    """

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for i in range(max_retries):
                try:
                    func(*args, **kwargs)
                except AssertionError as e:
                    if i == max_retries - 1:
                        raise e
                    continue
                break

        return wrapper

    return decorator


from typing import Dict, List, Union

from generator import model

from .dotnet_commons import TypeData
from .dotnet_constants import NAMESPACE
from .dotnet_helpers import class_wrapper, get_usings, namespace_wrapper

SPECIAL_CLASSES = [
    "LSPObject",
    "LSPAny",
    "LSPArray",
    "ChangeAnnotationIdentifier",
    "Pattern",
    "DocumentSelector",
    "InitializedParams",
]

def generate_special_classes(spec: model.LSPModel, types: TypeData) -> None:
    """Generate code for special classes in LSP."""
    for special_class in SPECIAL_CLASSES:
        for class_def in spec.structures + spec.typeAliases:
            if class_def.name == special_class:
                generate_special_class(class_def, spec, types)

def generate_special_class(
    type_def: Union[model.Structure, model.TypeAlias],
    spec: model.LSPModel,
    types: TypeData,
) -> Dict[str, str]:
    """Generate code for a special class."""
    lines: List[str] = []

    if type_def.name == "LSPObject":
        lines = namespace_wrapper(
            NAMESPACE,
            get_usings(["Dictionary", "DataContract", "JsonConverter"]),
            class_wrapper(
                type_def,
                ["public LSPObject(Dictionary<string, object?> value):base(value){}"],
                "Dictionary<string, object?>",
                ["[JsonConverter(typeof(CustomObjectConverter<LSPObject>))]"],
                is_record=False,
            ),
        )
    if type_def.name == "InitializedParams":
        lines = namespace_wrapper(
            NAMESPACE,
            get_usings(["Dictionary", "DataContract", "JsonConverter"]),
            class_wrapper(
                type_def,
                [
                    "public InitializedParams(Dictionary<string, object?> value):base(value){}"
                ],
                "Dictionary<string, object?>",
                ["[JsonConverter(typeof(CustomObjectConverter<InitializedParams>))]"],
                is_record=False,
            ),
        )
    if type_def.name == "LSPAny":
        lines = namespace_wrapper(
            NAMESPACE,
            get_usings(["DataContract", "JsonConverter"]),
            class_wrapper(
                type_def,
                [
                    "public LSPAny(object? value){this.Value = value;}",
                    "public object? Value { get; set; }",
                ],
                "object",
                ["[JsonConverter(typeof(LSPAnyConverter))]"],
            ),
        )
    if type_def.name == "LSPArray":
        lines = namespace_wrapper(
            NAMESPACE,
            get_usings(["DataContract", "List"]),
            class_wrapper(
                type_def,
                ["public LSPArray(List<object> value):base(value){}"],
                "List<object>",
                is_record=False,
            ),
        )

    if type_def.name == "Pattern":
        inner = [
            "private string pattern;",
            "public Pattern(string value){pattern = value;}",
            "public static implicit operator Pattern(string value) => new Pattern(value);",
            "public static implicit operator string(Pattern pattern) => pattern.pattern;",
            "public override string ToString() => pattern;",
        ]
        lines = namespace_wrapper(
            NAMESPACE,
            get_usings(["JsonConverter", "DataContract"]),
            class_wrapper(
                type_def,
                inner,
                None,
                [f"[JsonConverter(typeof(CustomStringConverter<{type_def.name}>))]"],
            ),
        )

    if type_def.name == "ChangeAnnotationIdentifier":
        inner = [
            "private string identifier;",
            "public ChangeAnnotationIdentifier(string value){identifier = value;}",
            "public static implicit operator ChangeAnnotationIdentifier(string value) => new ChangeAnnotationIdentifier(value);",
            "public static implicit operator string(ChangeAnnotationIdentifier identifier) => identifier.identifier;",
            "public override string ToString() => identifier;",
        ]
        lines = namespace_wrapper(
            NAMESPACE,
            get_usings(["JsonConverter", "DataContract"]),
            class_wrapper(
                type_def,
                inner,
                None,
                [f"[JsonConverter(typeof(CustomStringConverter<{type_def.name}>))]"],
            ),
        )

    if type_def.name == "DocumentSelector":
        inner = [
            "private DocumentFilter[] Filters { get; set; }",
            "public DocumentSelector(params DocumentFilter[] filters)",
            "{",
            "    Filters = filters ?? Array.Empty<DocumentFilter>();",
            "}",
            "public DocumentFilter this[int index]",
            "{",
            "    get { return Filters[index]; }",
            "    set { Filters[index] = value; }",
            "}",
            "public int Length => Filters.Length;",
            "public static implicit operator DocumentSelector(DocumentFilter[] filters) => new(filters);",
            "public static implicit operator DocumentFilter[](DocumentSelector selector) => selector.Filters;",
            "public IEnumerator<DocumentFilter> GetEnumerator() => ((IEnumerable<DocumentFilter>)Filters).GetEnumerator();",
            "System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator() => Filters.GetEnumerator();",
        ]
        lines = namespace_wrapper(
            NAMESPACE,
            get_usings(["JsonConverter", "DataContract"]),
            class_wrapper(
                type_def,
                inner,
                "IEnumerable<DocumentFilter>",
                [f"[JsonConverter(typeof(DocumentSelectorConverter))]"],
            ),
        )

    types.add_type_info(type_def, type_def.name, lines)