""" Main module for testing optimizers """
# Load libraries and pick the CUDA device if available
import json
import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
from torch.utils.tensorboard import SummaryWriter
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np
import time
# Custom libraries

from CustomOptimizer import *

from src.attack import FastGradientSignUntargeted
#from LinfPGDAttack import *
#import torchattacks

from loader import *
from src.models import MODELS_MAP
from src.utils import *

def flat_weight_dump(model):
    """ Returns a 1-d tensor containing all the network weights """
    is_empty = True
    for _, param in model.named_parameters():
        if param.requires_grad:
            if is_empty:
                flat_tensor = param.data.flatten()
                is_empty = False
            else:
                flat_tensor = torch.cat([flat_tensor, param.data.flatten()])
    return flat_tensor

def test(model, loader, adv_test=False, use_pseudo_label=False):
        # adv_test is False, return adv_acc as -1

        total_acc = 0.0
        num = 0
        total_adv_acc = 0.0

        with torch.no_grad():
            for data, label in loader:
                data, label = tensor2cuda(data), tensor2cuda(label)

                output = model(data, _eval=True)

                pred = torch.max(output, dim=1)[1]
                te_acc = evaluate(pred.cpu().numpy(), label.cpu().numpy(), 'sum')

                total_acc += te_acc
                num += output.shape[0]

                if adv_test:
                    # use predicted label as target label
                    with torch.enable_grad():
                        adv_data = attack.perturb(data,
                                                       pred if use_pseudo_label else label,
                                                       'mean',
                                                       False)

                    adv_output = model(adv_data, _eval=True)

                    adv_pred = torch.max(adv_output, dim=1)[1]
                    adv_acc = evaluate(adv_pred.cpu().numpy(), label.cpu().numpy(), 'sum')
                    total_adv_acc += adv_acc
                else:
                    total_adv_acc = -num

        return total_acc / num , total_adv_acc / num

# Trains the network
def train(model,tr_loader,va_loader=None, adv_train=False):

        opt = optimizer
        scheduler = torch.optim.lr_scheduler.MultiStepLR(opt,
                                                         milestones=[40000, 60000],
                                                         gamma=0.1)
        _iter = 0

        begin_time = time.time()

        writer1 = SummaryWriter(config_tb_path_train)
        writer2 = SummaryWriter(config_tb_path_test)

        for epoch in range(1, config_epochs+1):
            for data, label in tr_loader:
                data, label = tensor2cuda(data), tensor2cuda(label)

                if adv_train:
                    # When training, the adversarial example is created from a random
                    # close point to the original data point. If in evaluation mode,
                    # just start from the original data point.
                    adv_data = attack.perturb(data, label, 'mean', True)
                    output = model(adv_data, _eval=False)
                else:
                    output = model(data, _eval=False)

                loss = F.cross_entropy(output, label)

                opt.zero_grad()
                loss.backward()
                opt.step()

                if _iter % config_n_eval_step == 0:
                    t1 = time.time()

                    if adv_train:
                        with torch.no_grad():
                            stand_output = model(data, _eval=True)
                        pred = torch.max(stand_output, dim=1)[1]

                        # print(pred)
                        std_acc = evaluate(pred.cpu().numpy(), label.cpu().numpy()) * 100

                        pred = torch.max(output, dim=1)[1]
                        # print(pred)
                        adv_acc = evaluate(pred.cpu().numpy(), label.cpu().numpy()) * 100

                    else:

                        adv_data = attack.perturb(data, label, 'mean', False)

                        with torch.no_grad():
                            adv_output = model(adv_data, _eval=True)
                        pred = torch.max(adv_output, dim=1)[1]
                        # print(label)
                        # print(pred)
                        adv_acc = evaluate(pred.cpu().numpy(), label.cpu().numpy()) * 100

                        pred = torch.max(output, dim=1)[1]
                        # print(pred)
                        std_acc = evaluate(pred.cpu().numpy(), label.cpu().numpy()) * 100

                    t2 = time.time()

                    logger.info(f'epoch: {epoch}, iter: {_iter}, lr={opt.param_groups[0]["lr"]}, '
                                f'spent {time.time()-begin_time:.2f} s, tr_loss: {loss.item():.3f}')

                    logger.info(f'standard acc: {std_acc:.3f}%, robustness acc: {adv_acc:.3f}%')

                    writer1.add_scalar('Accuracy', std_acc, _iter)

                    writer1.add_scalar('Adversarial Accuracy', adv_acc, _iter)
                    # begin_time = time()

                    # if va_loader is not None:
                    #     va_acc, va_adv_acc = self.test(model, va_loader, True)
                    #     va_acc, va_adv_acc = va_acc * 100.0, va_adv_acc * 100.0

                    #     logger.info('\n' + '='*30 + ' evaluation ' + '='*30)
                    #     logger.info('test acc: %.3f %%, test adv acc: %.3f %%, spent: %.3f' % (
                    #         va_acc, va_adv_acc, time() - begin_time))
                    #     logger.info('='*28 + ' end of evaluation ' + '='*28 + '\n')

                    begin_time = time.time()

                _iter += 1
                # scheduler depends on training interation
                scheduler.step()

            if va_loader is not None:
                t1 = time.time()
                va_acc, va_adv_acc = test(model, va_loader, True, False)
                va_acc, va_adv_acc = va_acc * 100.0, va_adv_acc * 100.0

                t2 = time.time()
                logger.info('\n'+'='*20 +f' evaluation at epoch: {epoch} iteration: {_iter} ' \
                    +'='*20)
                logger.info(f'test acc: {va_acc:.3f}%, test adv acc: {va_adv_acc:.3f}%, spent: {t2-t1:.3f} s')
                logger.info('='*28+' end of evaluation '+'='*28+'\n')

                writer2.add_scalar('Accuracy', va_acc, _iter)

                writer2.add_scalar('Adversarial Accuracy', va_adv_acc, _iter)

        writer2.close()
        writer1.close()

# ################
parser = argparse.ArgumentParser(description='PyTorch Adversarial Training')
parser.add_argument(
    '--config', default='config.json', type=str, help='config file')
parser.add_argument('--resume', '-r', action='store_true',
                    help='resume from checkpoint')
args = parser.parse_args()

with open(args.config) as config_file:
    config = json.load(config_file)
config_experiment_number = config['experiment_number']
config_dataset = config['dataset']
config_architecture = config['architecture']
config_batch_size = config['batch_size']
config_optimizer = config['optimizer']
config_lr = config['lr']
config_momentum = config['momentum']
config_prune_epoch = config["prune_epoch"]
config_unfreeze_epoch = config["unfreeze_epoch"]
config_perc_to_prune = config['perc_to_prune']
config_step_of_prune = config["step_of_prune"]
config_radius = config['radius']
config_epochs = config['epochs']
config_tb_path_test = config['tb_path_test']
config_tb_path_train = config['tb_path_train']
config_batch_statistics_freq = config['batch_statistics_freq']
config_dump_movement = bool(config['dump_movement'] == 1)
config_projected = bool(config['projected'] == 1)
config_weight_decay = config['weight_decay']
config_radius = config['radius']
config_random_seed = config['random_seed']
config_gamma0 = config['gamma0']

config_one_shot_prune = config["one_shot_prune"]
config_iterative_prune = config["iterative_prune"]
config_epochs_to_finetune = config["epochs_to_finetune"]
config_epochs_to_densetrain = config["epochs_to_densetrain"]
#config_initial_accumulator_value = config['initial_accumulator_value']
#config_beta = config['beta']
#config_eps = config['eps']

config_n_eval_step = config['n_eval_step']
config_log_folder = config['log_folder']
config_epsilon = config['epsilon']
config_k = config['k']
config_alpha = config['alpha']

use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
#Set random seed
torch.manual_seed(config_random_seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(config_random_seed)

# Load data
if config_dataset == 'MNIST':
    trainloader, testloader = mnist_loader(batch_size=config_batch_size)
else:
    trainloader, testloader = cifar_loader(batch_size=config_batch_size)

model = MODELS_MAP[config_architecture]()
net = model.to(device)

criterion = nn.CrossEntropyLoss()

if config_optimizer == 0:
    optimizer = optim.SGD(
      net.parameters(), lr=config_lr,
      momentum=config_momentum, weight_decay=config_weight_decay)
elif config_optimizer == 1:
    optimizer = optim.Adagrad(
      net.parameters(), lr=config_lr, weight_decay=config_weight_decay)
elif config_optimizer == 2:
    optimizer = optim.Adam(net.parameters(), lr=config_lr, amsgrad=0, weight_decay=config_weight_decay)
elif config_optimizer == 3:
    optimizer = optim.Adam(net.parameters(), lr=config_lr, amsgrad=1, weight_decay=config_weight_decay)
elif config_optimizer == 4:
    optimizer = optim.RMSprop(net.parameters(), lr=config_lr)
elif config_optimizer == 5:
    optimizer = AdaptiveLinearCoupling(
        net.parameters(), lr=config_lr,
        weight_decay=config_weight_decay)
elif config_optimizer == 6:
    #optimizer = AdaACSA(
    #    net.parameters(), lr=config_lr, radius=1, projected=config_projected)
    optimizer = AdaACSA(
        net.parameters(), lr=config_lr, radius=config_radius,
        weight_decay=config_weight_decay, projected=config_projected,
        gamma0=config_gamma0, beta=config_beta,
        initial_accumulator_value=config_initial_accumulator_value,
        eps=config_eps)
elif config_optimizer == 7:
    optimizer = AdaAGDplus(
        net.parameters(), lr=config_lr, radius=config_radius, projected=config_projected,
        initial_accumulator_value=config_initial_accumulator_value,
        eps=config_eps)
elif config_optimizer == 8:
    optimizer = AdaJRGS(
        net.parameters(), lr=config_lr, radius=config_radius, projected=config_projected,
        initial_accumulator_value=config_initial_accumulator_value,
        eps=config_eps)
elif config_optimizer == 9:
    optimizer = CustomOptimizer(net.parameters(),lr=config_lr,
    momentum=config_momentum,
    weight_decay=config_weight_decay,
    len_step = len(trainloader),

    one_shot_prune  = config_one_shot_prune,
    prune_epoch=config_prune_epoch,
    step_of_prune=config_step_of_prune,
    perc_to_prune = config_perc_to_prune,

    iterative_prune = config_iterative_prune,
    unfreeze_epoch=config_unfreeze_epoch,
    epochs_to_densetrain = config_epochs_to_densetrain,
    epochs_to_finetune= config_epochs_to_finetune
   )

# Writer path for display on TensorBoard
if not os.path.exists(config_tb_path_test):
    os.makedirs(config_tb_path_test)
if not os.path.exists(config_tb_path_train):
    os.makedirs(config_tb_path_train)

if not os.path.exists(config_log_folder):
    os.makedirs(config_log_folder)
#path_name = config_tb_path + \
#    str(config_experiment_number) + "_" + str(optimizer)

logger = create_logger(config_epsilon,config_log_folder,'train','info')

# Initialize weights
net.apply(weights_init_uniform_rule)

#load the model
'''
ckp_path = 'checkpoints/epoch_model_40.pth'
checkpoint_model, start_epoch = load_ckp(ckp_path, net)
'''
attack = FastGradientSignUntargeted(net,
                                        config_epsilon,
                                        config_alpha,
                                        min_val=0,
                                        max_val=1,
                                        max_iters=config_k,
                                        _type='linf')

train(net,trainloader,testloader,adv_train=True)

# Dump some info on the range of parameters after training is finished
for param in net.parameters():
    print(str(torch.min(param.data).item()) + " " + str(torch.max(param.data).item()))


#! /usr/bin/env python3
"""
 # @ Author: Kenneth Simon
 # @ Email: smkk00715@gmail.com
 # @ Create Time: 2022-05-14 15:46:24
 # @ Modified time: 2022-05-14 16:01:51
 # @ Description:
 """

import csv
import datetime
import json
import os
import sys
import time

import ipdb
import numpy as np
import psutil
import rosparam
import rospy
import sympy
import tf.transformations as tft
from gazebo_msgs.msg import ModelStates
from geometry_msgs.msg import Pose, PoseStamped, Twist, TwistStamped
from matplotlib import pyplot as plt
from sensor_msgs.msg import Imu
from std_msgs.msg import Bool, Float32, Float32MultiArray

from utils import *

class UAVStateMonitor:
    """get the realtime state of the uav"""

    def __init__(self, uav_name) -> None:
        self.uav_name = uav_name
        self.real_pose_sub = rospy.Subscriber(
            f"/{self.uav_name}/real_pose", PoseStamped, self.real_pose_cb, queue_size=1
        )
        self.real_twist_sub = rospy.Subscriber(
            f"/{self.uav_name}/real_twist", TwistStamped, self.real_twist_cb, queue_size=1
        )
        self.curr_pose = PoseStamped()
        self.curr_twist = TwistStamped()

    def real_pose_cb(self, real_pose_msg: PoseStamped):
        self.curr_pose = real_pose_msg

    def real_twist_cb(self, real_twist_msg: TwistStamped):
        self.curr_twist = real_twist_msg

class PowerAnalyzer:
    """PowerAnalyzer for one uav.

    Refer to:
    Z. Liu, R. Sengupta and A. Kurzhanskiy, "A power consumption model for multi-rotor small
    unmanned aircraft systems," 2017 International Conference on Unmanned Aircraft Systems (ICUAS),
    Miami, FL, USA, 2017, pp. 310-315, doi: 10.1109/ICUAS.2017.7991310.
    """

    # Parameters for iris quadrotor
    N = 2  # totol number of blades in a single propeller
    A_quad = 0.03891  # cross-sectional area of the uav when against wind (cm^2)
    c = 1  # blade chord width
    c_d = 0.5  # air drag coefficient of the blade. a normal car's drag coefficient is 0.28-0.4
    k_1 = 0.8554
    k_2 = 0.3051  # kg/m^(1/2)
    k_3 = 8.54858e-06  # motor constant (N/(rad/s)^2)
    c_2 = 0.3177  # m/kg^(1/2)
    c_3 = 0  # negligible, approx 0
    c_4 = 0.0296  # kg/m
    c_5 = 0.0279  # Ns/m
    c_6 = 0  # negligible, approx 0
    g = 9.8
    pho = 1.225  # kg/m^3 air density
    """
    Size: 10 cm in height, 55 cm motor-to-motor, 0.3891m in side length
    Propellers: (2) 10 x 4.7 normal-rotation, (2) 10 x 4.7 reverse-rotation
    Motors: AC 2830, 850 kV
    Battery: 3-cell 11.1 V 3.5 Ah lithium polymer with XT-60 type connector. Weight: 262 grams
    Weight (with battery): 1282 grams
    Motor to motor dimension: 550mm
    """

    def __init__(self, uav_name="iris_0", state_monitor=None) -> None:
        self.uav_name = uav_name
        self.thrust_est = 0  # estimated thrust
        self.instant_power = 0
        self.energy_total = 0
        self.induced_power = 0
        self.profile_power = 0
        self.parasite_power = 0
        self.total_induced_power = 0
        self.total_profile_power = 0
        self.total_parasite_power = 0
        self.total_instant_power = 0
        self.sampling_rate = 10  # (hz)

        # subs
        self.thrust_est_sub = rospy.Subscriber(f"/{self.uav_name}/thrust_est", Float32, self.thrust_est_cb)
        if state_monitor is None:
            self.state_monitor = UAVStateMonitor(uav_name=self.uav_name)
        else:
            self.state_monitor = state_monitor
        # self.model_state_sub = rospy.Subscriber(f"/gazebo/model_states", ModelStates, self.model_state_cb)
        # self.imu_sub = rospy.Subscriber(f"/{self.uav_name}/imu/data", Imu, self.imu_cb)

        # pubs
        self.induced_power_pub = rospy.Publisher(f"/{self.uav_name}/perf_analyzer/induced_power", Float32, queue_size=1)
        self.profile_power_pub = rospy.Publisher(f"/{self.uav_name}/perf_analyzer/profile_power", Float32, queue_size=1)
        self.parasite_power_pub = rospy.Publisher(
            f"/{self.uav_name}/perf_analyzer/parasite_power", Float32, queue_size=1
        )
        self.instant_power_pub = rospy.Publisher(f"/{self.uav_name}/perf_analyzer/instant_power", Float32, queue_size=1)

    def calcu_induced_power(self, T, V_vert):
        """
        Calculate the induced power.

        Parameters:
        - T: Total thrust.
        - V_vert: Vertical speed at each boundary.

        Returns:
        Induced power based on given thrust and vertical speed.
        """
        return PowerAnalyzer.k_1 * T * (V_vert / 2 + np.sqrt((V_vert / 2) ** 2 + T / PowerAnalyzer.k_2**2))

    def calcu_profile_power(self, T, V_air, alpha):
        """Parameters:
        - T: Total thrust.
        - V_air: Horizontal airspeed.
        - alpha: angle of attack for each propeller disk. alpha = pi/2 - pitch

        """
        return PowerAnalyzer.c_2 * T ** (3 / 2) + PowerAnalyzer.c_3 * (V_air * np.cos(alpha)) ** 2 * T ** (1 / 2)

    def calcu_parasite_power(self, C_d, rho, A_quad, V_air):
        """
        Parameters:
        - C_d: Drag coefficient of the vehicle body.
        - rho: Density of air.
        - A_quad: Cross-sectional area of the vehicle when against wind.
        - V_air: Horizontal airspeed.
        """
        return 0.5 * C_d * rho * A_quad * V_air**3

    def calculate_T(m, V_air, alpha):
        """
        the theoretical thrust given the airspeed and angle of attack.

        - alpha: Angle of attack of the propeller disk
        """
        T = sympy.symbols("T")
        f = (
            np.sqrt(
                (m * PowerAnalyzer.g - (PowerAnalyzer.c_5 * (V_air * np.cos(alpha)) ** 2 + PowerAnalyzer.c_6 * T)) ** 2
                + (PowerAnalyzer.c_4 * V_air**2) ** 2
            )
            - T
        )
        result = sympy.solve(f, T)
        return result

    def test_power_model(self):
        m = 1.5
        ax1 = plt.subplot(311)
        ax2 = plt.subplot(312)
        ax3 = plt.subplot(313)

        # induced power
        T = np.arange(0, 10, 0.1)
        # V_vert = np.arange(0, 10, 0.1)
        self.induced_power = self.calcu_induced_power(T, 0)
        ax1.plot(T, self.induced_power)
        ax1.set_xlabel("Thrust (N)")
        ax1.set_ylabel("Induced power (W)")
        ax1.set_title("Induced power vs thrust")
        ax1.legend()

        # profile power
        T = np.arange(0, 10, 0.1)
        V_air = 10
        alpha = np.radians(10)
        self.profile_power = self.calcu_profile_power(T, V_air, alpha)
        ax2.plot(
            T,
            self.profile_power,
        )
        ax2.set_xlabel("Thrust (N)")
        ax2.set_ylabel("Profile power (W)")
        ax2.set_title("Profile power vs thrust and v_air")
        ax2.legend()

        # # parasite power
        V_air = np.arange(0, 10, 0.1)
        self.parasite_power = self.calcu_parasite_power(self.c_d, self.pho, self.A_quad, V_air)
        ax3.plot(V_air, self.parasite_power)
        ax3.set_xlabel("V_air (m/s)")
        ax3.set_ylabel("Parasite power (W)")
        ax3.set_title("Parasite power vs V_air")
        ax3.legend()

        plt.legend()
        plt.tight_layout()
        plt.show()

    def thrust_est_cb(self, thrust_est_msg: Float32):
        self.thrust_est = thrust_est_msg.data

    def model_state_cb(self, model_state_msg: ModelStates):
        model_idx = model_state_msg.name.index(f"{self.uav_name}")
        self.twist = model_state_msg.twist[model_idx]
        self.pose = model_state_msg.pose[model_idx]

    def run(self):
        loop_rate = rospy.Rate(self.sampling_rate)
        rospy.loginfo(f"Power analyzer for uav {self.uav_id} started.")
        while not rospy.is_shutdown():
            self.pose = self.state_monitor.curr_pose.pose
            self.twist = self.state_monitor.curr_twist.twist
            [r, p, y] = tft.euler_from_quaternion(
                [self.pose.orientation.x, self.pose.orientation.y, self.pose.orientation.z, self.pose.orientation.w]
            )
            [r, p, y, _, _, _] = ENU2NED_Position([r, p, y, 0, 0, 0])
            speed_in_body = inertial2body([self.twist.linear.x, self.twist.linear.y, self.twist.linear.z], self.pose)

            self.induced_power = self.calcu_induced_power(self, self.thrust_est, abs(speed_in_body[2]))
            self.profile_power = self.calcu_profile_power(self.thrust_est, abs(speed_in_body[0]), p)
            self.parasite_power = self.calcu_parasite_power(self.c_d, self.pho, self.A_quad, abs(speed_in_body[0]))
            self.instant_power = self.induced_power + self.profile_power + self.parasite_power
            self.energy_total += self.instant_power / self.sampling_rate

            self.induced_power_pub.publish(self.induced_power)
            self.profile_power_pub.publish(self.profile_power)
            self.parasite_power_pub.publish(self.parasite_power)
            self.instant_power_pub.publish(self.instant_power)
            loop_rate.sleep()

    def run_in_ext(self, sampling_duty):
        """run in external loop"""
        self.pose = self.state_monitor.curr_pose.pose
        self.twist = self.state_monitor.curr_twist.twist
        [r, p, y] = tft.euler_from_quaternion(
            [self.pose.orientation.x, self.pose.orientation.y, self.pose.orientation.z, self.pose.orientation.w]
        )
        [r, p, y, _, _, _] = ENU2NED_Position([r, p, y, 0, 0, 0])
        speed_in_body = inertial2body([self.twist.linear.x, self.twist.linear.y, self.twist.linear.z], self.pose)

        if isinstance(self.thrust_est, complex) or np.isnan(self.thrust_est):
            rospy.logerr(f"thrust_est: {self.thrust_est}")
            return
        if isinstance(speed_in_body[2], complex) or np.isnan(speed_in_body[2]):
            rospy.logerr(f"speed_in_body[2]: {speed_in_body[2]}")
            return
        self.induced_power = self.calcu_induced_power(self.thrust_est, abs(speed_in_body[2]))
        if isinstance(self.induced_power, complex) or np.isnan(self.induced_power):
            rospy.logerr(
                f"induced_power: {self.induced_power} thrust_est: {self.thrust_est} speed_in_body[2]: {speed_in_body[2]}"
            )
            return
        self.profile_power = self.calcu_profile_power(self.thrust_est, abs(speed_in_body[0]), p)
        self.parasite_power = self.calcu_parasite_power(self.c_d, self.pho, self.A_quad, abs(speed_in_body[0]))
        self.total_induced_power += self.induced_power
        self.total_profile_power += self.profile_power
        self.total_parasite_power += self.parasite_power
        self.instant_power = self.induced_power + self.profile_power + self.parasite_power
        self.total_instant_power += self.instant_power
        self.energy_total += self.instant_power * sampling_duty

        self.induced_power_pub.publish(self.induced_power)
        self.profile_power_pub.publish(self.profile_power)
        self.parasite_power_pub.publish(self.parasite_power)
        self.instant_power_pub.publish(self.instant_power)

class NetworkAnalyzer:
    def __init__(self, uav_name="iris_0", state_monitor=None) -> None:
        self.uav_name = uav_name
        self.tgt_pose = PoseStamped()
        self.net_cap_th = 3  # (mbps) network capacity reach threshold
        self.net_cap_curr = 0  # (mbps) current network capacity
        self.net_cap_tgt = 0
        self.total_traffic = 0  # mb
        self.cap_resp_time = None  # s reponse time from the timestamp not meeting the capacity to meeting
        self.cap_resp_time_total = 0  # s
        self.cap_resp_time_total_weighted = 0  # s
        self.cap_resp_time_ave = 0  # s average time to meet the capacity
        self.timestamp_not_meet_cap = None  # s timestamp when not meeting the capacity requirement
        self.cap_meet_flag = False
        self.start_time = None
        self.cap_grad_curr = 0
        self.cap_grad_in_win = 0  # traffic gradient in a time window

        # params for wifi channel
        self.band_width = 2.4e9  # 2.4Ghz
        self.transmit_power = dbm2watt(20)  # dbm
        self.noise_power = dbm2watt(25)

        if state_monitor is None:
            self.state_monitor = UAVStateMonitor(uav_name=self.uav_name)
        else:
            self.state_monitor = state_monitor

        # subs
        self.tgt_pose_sub = rospy.Subscriber(
            f"/{self.uav_name}/target_pose", PoseStamped, self.tgt_pose_cb, queue_size=1
        )
        self.cap_tgt_sub = rospy.Subscriber(f"/crowds/target_net_capacity", Float32MultiArray, self.cap_tgt_cb)
        self.cap_tgt_grad_sub = rospy.Subscriber(f"/scheduler/cap_tgt_grad", Float32MultiArray, self.cap_tgt_grad_cb)
        # pubs
        self.curr_cap_pub = rospy.Publisher(f"/{self.uav_name}/perf_analyzer/curr_net_cap", Float32, queue_size=1)
        self.test_data_pub = rospy.Publisher(f"/{self.uav_name}/perf_analyzer/test", Float32MultiArray, queue_size=1)

    def cap_tgt_grad_cb(self, msg: Float32MultiArray):
        idx = int(self.uav_name.split("_")[-1])
        self.cap_grad_curr = msg.data[idx]

    def tgt_pose_cb(self, tgt_pose_msg: PoseStamped):
        self.tgt_pose = tgt_pose_msg

    def cap_tgt_cb(self, tgt_cap_msg: Float32):
        idx = int(self.uav_name.split("_")[-1])
        self.net_cap_tgt = tgt_cap_msg.data[idx]

    def run_in_ext(self, samping_duty):
        dist = np.linalg.norm(
            [
                self.tgt_pose.pose.position.x - self.state_monitor.curr_pose.pose.position.x,
                self.tgt_pose.pose.position.y - self.state_monitor.curr_pose.pose.position.y,
                0 - self.state_monitor.curr_pose.pose.position.z,
            ]
        )
        path_loss = two_ray_ground_reflection(dist, 1, self.state_monitor.curr_pose.pose.position.z, 1, 1)
        gain = 1 / path_loss
        self.net_cap_curr = calcu_capacity(self.band_width, self.transmit_power, gain, self.noise_power) / 1e6
        if self.net_cap_curr - self.net_cap_tgt >= -self.net_cap_th:  # meet the capacity requirement
            if (
                not self.cap_meet_flag and self.timestamp_not_meet_cap != None
            ):  # first time meet the capacity requirement
                self.cap_resp_time = rospy.get_time() - self.timestamp_not_meet_cap
                self.cap_resp_time_total += self.cap_resp_time
                # if self.uav_name == "iris_0":
                #     print(f"cap_resp_time: {self.cap_resp_time}")
                #     print(f"cap_resp_time_total: {self.cap_resp_time_total}")

                cap_resp_time_weighted = self.cap_resp_time * self.cap_grad_in_win
                self.cap_resp_time_total_weighted += cap_resp_time_weighted  # weighted by traffic gradient
                # if self.uav_name == "iris_0":
                #     print(f"traffic_grad_in_win: {self.cap_grad_in_win}")
                #     print(f"cap_resp_time_weighted: {cap_resp_time_weighted}")
                #     print(f"cap_resp_time_total_weighted: {self.cap_resp_time_total_weighted}")
                #     print("-" * 50)
                self.cap_grad_in_win = 0
            self.cap_meet_flag = True
        else:
            if self.timestamp_not_meet_cap is None:  # initial timestamp
                self.timestamp_not_meet_cap = rospy.get_time()
            if self.cap_meet_flag:  # change-point: last loop meet while this loop not meet
                self.timestamp_not_meet_cap = rospy.get_time()
            self.cap_meet_flag = False
            if self.cap_grad_curr > 0:
                self.cap_grad_in_win += self.cap_grad_curr
        self.total_traffic += self.net_cap_curr * samping_duty

        self.curr_cap_pub.publish(self.net_cap_curr)
        self.test_data_pub.publish(Float32MultiArray(data=[self.cap_grad_in_win]))

class PerformanceAnalyzer:
    def __init__(self, uav_num=3, samping_freq=10, controller="nmpc", mode="both", result_file_folder=None) -> None:
        self.uav_num = uav_num
        self.samping_freq = samping_freq  # (hz)
        self.uav_state_monitors = []  # shared monitors for memory efficiency
        self.power_analyzers = []
        self.network_analyzers = []
        self.enable_simu_flag = False
        self.controller = controller
        self.csv_writers = []
        self.time_suffix = datetime.datetime.now().strftime("%m-%d-%H%M%S")
        self.ave_curr_net_cap = 0  # average current network capacity of all uavs
        self.ave_tgt_net_cap = 0
        self.ave_instant_power = 0
        self.perf_data_write_freq = 2  # hz
        self.cpu_usage = 0
        self.cpu_usage_sum = 0
        self.cpu_usage_ave = 0
        if result_file_folder is None:
            self.result_file_folder = os.path.join(sys.path[0], f"../data/perf_anaysis")
        else:
            self.result_file_folder = result_file_folder
        if mode == "both":
            self.mode = 3
        elif mode == "power":
            self.mode = 1
            self.enable_simu_flag = True
        elif mode == "net":
            self.mode = 2
        else:
            raise ValueError("mode must be both, power or net")
        print(f"{'uav_num':<20}: {self.uav_num}")
        print(f"{'controller':<20}: {self.controller}")
        print(f"{'mode':<20}: {self.mode}")
        print(f"{'sampling_freq':<20}: {self.samping_freq}")
        print(f"{'result_file_folder':<20}: {self.result_file_folder}")

        if not os.path.exists(self.result_file_folder):
            os.makedirs(self.result_file_folder)
        for i in range(uav_num):
            self.uav_state_monitors.append(UAVStateMonitor(uav_name=f"iris_{i}"))
            if self.mode != 2:  # power or both
                self.power_analyzers.append(
                    PowerAnalyzer(uav_name=f"iris_{i}", state_monitor=self.uav_state_monitors[i])
                )
            if self.mode >= 2:  # net or both
                self.network_analyzers.append(
                    NetworkAnalyzer(uav_name=f"iris_{i}", state_monitor=self.uav_state_monitors[i])
                )
            self.csv_writers.append(
                csv.writer(
                    open(
                        os.path.join(
                            sys.path[0],
                            self.result_file_folder,
                            f"perf_data_{controller}_u{uav_num}_iris_{i}_{self.time_suffix}.csv",
                        ),
                        "w",
                    )
                )
            )
            self.csv_writers[i].writerow(  # header
                [
                    "time",
                    "induced_power",
                    "profile_power",
                    "parasite_power",
                    "instant_power",
                    "energy_total",
                    "net_cap_curr",
                    "net_cap_tgt",
                    "total_traffic",
                    "cap_resp_time",
                    "cap_resp_time_total_weighted",
                ]
            )

        self.enable_simu_sub = rospy.Subscriber("/enable_simu", Bool, self.start_simu_cb, queue_size=1)

        # pubs
        self.ave_curr_net_cap_pub = rospy.Publisher(f"/perf_analyzer/ave_curr_net_cap", Float32, queue_size=1)
        self.ave_tgt_net_cap_pub = rospy.Publisher(f"/perf_analyzer/ave_tgt_net_cap", Float32, queue_size=1)
        self.ave_instant_power_pub = rospy.Publisher(f"/perf_analyzer/ave_instant_power", Float32, queue_size=1)
        self.cpu_usage_pub = rospy.Publisher(f"/perf_analyzer/cpu_usage", Float32, queue_size=1)
        self.test_pub = rospy.Publisher(f"/perf_analyzer/test", Float32MultiArray, queue_size=1)

    def start_simu_cb(self, start_simu_msg: Bool):
        self.enable_simu_flag = start_simu_msg.data

    def run(self):
        loop_rate = rospy.Rate(self.samping_freq)
        while not self.enable_simu_flag and not rospy.is_shutdown():
            loop_rate.sleep()
            print(f"[{rospy.get_time():<20}] waiting for simulation to start...", end="\r")
            continue
        rospy.loginfo("Performance analyzer started.")
        start_time = rospy.get_time()
        cnt = 0
        while not rospy.is_shutdown():
            if not self.enable_simu_flag:
                break
            for i in range(self.uav_num):
                if self.mode != 2:
                    self.power_analyzers[i].run_in_ext(1 / self.samping_freq)
                if self.mode >= 2:
                    self.network_analyzers[i].run_in_ext(1 / self.samping_freq)
                if cnt % (int(self.samping_freq / self.perf_data_write_freq)) == 0:
                    if self.mode == 1:
                        self.curr_perf_data = [
                            rospy.get_time(),
                            self.power_analyzers[i].induced_power,
                            self.power_analyzers[i].profile_power,
                            self.power_analyzers[i].parasite_power,
                            self.power_analyzers[i].instant_power,
                            self.power_analyzers[i].energy_total,
                            0,
                            0,
                            0,
                            0,
                            0,
                        ]
                    elif self.mode == 2:
                        self.curr_perf_data = [
                            0,
                            0,
                            0,
                            0,
                            0,
                            self.network_analyzers[i].net_cap_curr,
                            self.network_analyzers[i].net_cap_tgt,
                            self.network_analyzers[i].total_traffic,
                            self.network_analyzers[i].cap_resp_time,
                            self.network_analyzers[i].cap_resp_time_total_weighted,
                        ]
                        self.curr_perf_data.insert(0, rospy.get_time())
                    elif self.mode == 3:
                        self.curr_perf_data = [
                            rospy.get_time(),
                            self.power_analyzers[i].induced_power,
                            self.power_analyzers[i].profile_power,
                            self.power_analyzers[i].parasite_power,
                            self.power_analyzers[i].instant_power,
                            self.power_analyzers[i].energy_total,
                            self.network_analyzers[i].net_cap_curr,
                            self.network_analyzers[i].net_cap_tgt,
                            self.network_analyzers[i].total_traffic,
                            self.network_analyzers[i].cap_resp_time,
                            self.network_analyzers[i].cap_resp_time_total_weighted,
                        ]
                for data in self.curr_perf_data:
                    if data != None:
                        data = round(data, 4)
                self.csv_writers[i].writerow(self.curr_perf_data)
                self.ave_curr_net_cap += self.network_analyzers[i].net_cap_curr
                self.ave_tgt_net_cap += self.network_analyzers[i].net_cap_tgt
                self.ave_instant_power += self.power_analyzers[i].instant_power
            self.ave_curr_net_cap /= self.uav_num
            self.ave_tgt_net_cap /= self.uav_num
            self.ave_instant_power /= self.uav_num
            self.ave_curr_net_cap_pub.publish(self.ave_curr_net_cap)
            self.ave_tgt_net_cap_pub.publish(self.ave_tgt_net_cap)
            self.ave_instant_power_pub.publish(self.ave_instant_power)
            self.ave_curr_net_cap = 0
            self.ave_tgt_net_cap = 0
            self.ave_instant_power = 0
            self.cpu_usage = psutil.cpu_percent()
            self.cpu_usage_sum += self.cpu_usage
            self.cpu_usage_ave = self.cpu_usage_sum / (cnt + 1)
            self.cpu_usage_pub.publish(self.cpu_usage)
            test_msg = Float32MultiArray()
            test_msg.data = [self.cpu_usage_sum, self.cpu_usage_ave]  # FIXME
            self.test_pub.publish(test_msg)
            cnt += 1
            loop_rate.sleep()

        rospy.loginfo("Performance analyzer stopped.")
        end_time = rospy.get_time()
        file_path = os.path.join(
            sys.path[0],
            self.result_file_folder,
            f"perf_anaysis_result_{self.controller}_u{self.uav_num}_" + self.time_suffix + ".json",
        )
        file_obj = open(file_path, "w")
        result = {
            "cpu_usage_ave": self.cpu_usage_ave,
            "cpu_usage_sum": self.cpu_usage_sum,
        }
        for i in range(self.uav_num):
            # network performance
            net_traf_total = self.network_analyzers[i].total_traffic
            cap_resp_time_ave = self.network_analyzers[i].cap_resp_time_total / (end_time - start_time)
            cap_resp_time_total = self.network_analyzers[i].cap_resp_time_total
            cap_resp_time_total_weighted = self.network_analyzers[i].cap_resp_time_total_weighted
            cap_resp_time_ave_weighted = cap_resp_time_total_weighted / (end_time - start_time)
            # power performance
            induced_power_ave = self.power_analyzers[i].total_induced_power / (end_time - start_time)
            parasite_power_ave = self.power_analyzers[i].total_parasite_power / (end_time - start_time)
            profile_power_ave = self.power_analyzers[i].total_profile_power / (end_time - start_time)
            instant_power_ave = self.power_analyzers[i].total_instant_power / (end_time - start_time)
            energy_total = self.power_analyzers[i].energy_total

            result_dict = {
                "simu_time": end_time - start_time,
                "net_traf_total": net_traf_total,
                "cap_resp_time_ave": cap_resp_time_ave,
                "cap_resp_time_total": cap_resp_time_total,
                "cap_resp_time_total_weighted": cap_resp_time_total_weighted,
                "cap_resp_time_ave_weighted": cap_resp_time_ave_weighted,
                "induced_power_ave": induced_power_ave,
                "parasite_power_ave": parasite_power_ave,
                "profile_power_ave": profile_power_ave,
                "instant_power_ave": instant_power_ave,
                "energy_total": energy_total,
            }
            result[f"iris_{i}"] = result_dict
            uav_param_file = os.path.join(sys.path[0], f"../config/uav_params.yaml")
            uav_param_dict = rosparam.load_file(uav_param_file)[0][0]
            result[f"iris_{i}"]["param"] = uav_param_dict["nmpcc"]["power_saving"]

        json.dump(result, file_obj, indent=4)
        rospy.loginfo(f"result file saved to {file_path}")

if __name__ == "__main__":
    rospy.init_node("performace_analysis_node", anonymous=False)
    uav_num = rospy.get_param("~uav_num", 1)
    controller = rospy.get_param("~controller", "nmpc")
    mode = rospy.get_param("~mode", "both")
    result_file_folder = rospy.get_param("~result_file_folder", None)
    if result_file_folder == "default":
        result_file_folder = None

    # test PowerAnalyzer
    # analyzers = [PowerAnalyzer(uav_id=i) for i in range(uav_num)]
    # for i in range(uav_num):
    #     Thread(target=analyzers[i].run).start()

    # test NetworkAnalyzer
    # traffic_anas = [TrafficAnalyzer(uav_name=f"iris_{i}") for i in range(uav_num)]
    # freq = 30
    # loop_rate = rospy.Rate(freq)
    # while not rospy.is_shutdown():
    #     for i in range(uav_num):
    #         traffic_anas[i].run_in_ext(freq)
    #     loop_rate.sleep()
    # rospy.spin()

    perf_ana = PerformanceAnalyzer(
        uav_num=uav_num, controller=controller, mode=mode, result_file_folder=result_file_folder
    )
    perf_ana.run()


#!/usr/bin/env python
"""
fla.gr helper utils for mainly the redis db models

http://xkcd.com/353/

Josh Ashby
2013
http://joshashby.com
joshuaashby@joshashby.com
"""

def toBoolean(str):
    if str == 'True':
        return True
    elif str == 'False':
        return False
    else:
        raise Exception("Not a boolean")


import gym
from gym.spaces import *
# import gensim
import os
import numpy as np
import json
from xml.etree import ElementTree
from collections import Counter
import logging
import datetime
import sys
import math

class DStSNEEnv(gym.Env):
    metadata = {'render.modes': ['ansi']}
    segment_num = 20
    tsne_dim = 3

    def __init__(self):
        # placeholders
        self.topic_id = ""
        self.observation_space = Box(low=-10, high=10, shape=(20, DStSNEEnv92.segment_num, DStSNEEnv92.tsne_dim),
                                     dtype=float)
        self.action_space = Box(low=-1, high=1, shape=(DStSNEEnv92.tsne_dim,), dtype=float)
        # self.dct = gensim.corpora.Dictionary()
        # self.state = np.asarray([])  # tensor (#doc, #segment, 1)
        self.board = np.asarray([])  # tensor (#doc,  #segment, 3)
        self.board_backup = np.copy(self.board)
        # self.vocab_size = 0
        self.doc_num = 0
        self.doc_list = []
        # self.cupy_device = 0
        self.doc_rel = Counter()  # doc -> relevance score
        self.submitted = set()
        self.doc2idx = {}

        self.iter_cnt = 0
        self.acc_reward = 0
        self.logger = logging.getLogger(__name__)  # place holder, will be changed later
        self.name = self.__class__.__name__
        self.empty_value = np.random.random((DStSNEEnv92.segment_num, DStSNEEnv92.tsne_dim)) * 10
        # self.play_log = None
        # self.play = False

    def _set_logger(self, topic_id, env_rank):
        # file handler
        os.makedirs('ds_log/', exist_ok=True)
        fh = logging.FileHandler("ds_log/RL-{}-{}-{}-{}.log".format(
            self.name, topic_id, env_rank, datetime.datetime.now().strftime('%Y%m%d_%H%M%S_%f')))
        fh.setLevel(logging.DEBUG)

        # setup logger
        self.logger = logging.getLogger("{}_{}_{}".format(self.name, topic_id, env_rank))
        self.logger.setLevel(logging.DEBUG)
        # console handler
        ch = logging.StreamHandler(sys.stdout)
        ch.setLevel(logging.CRITICAL)
        # set format
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        fh.setFormatter(formatter)
        ch.setFormatter(formatter)

        self.logger.addHandler(fh)
        self.logger.addHandler(ch)
        self.logger.info("Logger Created")

    def set_topic(self, topic_id, env_rank, corpus_path, truth_path):
        self.topic_id = topic_id
        self._set_logger(topic_id, env_rank)

        self.logger.info("Running on Topic {}, Env Rank: {}".format(self.topic_id, env_rank))
        # self.cupy_device = gpu_id
        self.doc_list = json.load(open("data/small_corpus.json"))[topic_id]
        self.doc_num = len(self.doc_list)

        # set up doc2idx
        for idx, doc in enumerate(self.doc_list):
            self.doc2idx[doc] = idx

        # self.dct = gensim.corpora.Dictionary.load_from_text(os.path.join(corpus_path, topic_id, "dictionary.txt"))
        # self.vocab_size = len(self.dct)
        self.board = np.load(os.path.join(corpus_path, topic_id, "mat_norm_3.npz"))['tsne_norm']
        self.board = self.board.reshape((self.doc_num, DStSNEEnv92.segment_num, DStSNEEnv92.tsne_dim))
        self.board_backup = np.copy(self.board)

        self.observation_space = Box(low=np.min(self.board), high=np.max(self.board), shape=self.board.shape,
                                     dtype=float)
        self.action_space = Box(low=-1, high=1, shape=(DStSNEEnv92.tsne_dim,), dtype=float)

        # initial action

        # action = np.random.random((DStSNEEnv9.tsne_dim,))

        topic_data = ElementTree.parse(truth_path).find("./*/topic[@id=\"{}\"]".format(topic_id))

        for psg in topic_data.findall(".//passage"):
            doc = psg.find("./docno").text
            rel = int(psg.find("./rating").text)
            self.doc_rel[doc] += rel
        self.submitted.clear()
        self.iter_cnt = 0
        self.acc_reward = 0

        # self.play = play
        # if play:
        #     self.play_log = open(os.path.join("ds_log", "{}_{}_play.log".format(self.topic_id, env_rank)))

    def log_submit(self, submit_docs):
        log_str = ""
        for doc, score, rel in submit_docs:
            log_str += "{}\t{}\t{}\t{}\t{}\t\n".format(self.topic_id, self.iter_cnt, doc, score, rel)
        self.logger.info("\n" + log_str)

        # if self.play:
        #     self.play_log.write(log_str)

    def step(self, action):
        self.logger.info("Action: {}".format(action))

        rel = np.sum(np.dot(self.board, action), axis=1)

        doc_score = sorted(zip(self.doc_list, rel), key=lambda x: -x[1])

        submit = []
        reward = 0
        for doc, score in doc_score:
            if doc not in self.submitted:
                rel = self.doc_rel[doc]
            else:
                rel = 0
            # set selected documents to reserved values
            self.board[self.doc2idx[doc]] = self.empty_value
            submit.append((doc, score, rel))
            self.submitted.add(doc)
            reward += rel
            if len(submit) == 5:
                break

        observation = np.copy(self.board)

        self.log_submit(submit)
        self.acc_reward += reward
        self.iter_cnt += 1
        # half of the corpus are relevant
        done = (len(self.submitted) >= self.doc_num // 2)  or (self.iter_cnt >= 30)

        if done:
            self.logger.warning("Done at iteration {}, total reward {}\n".format(self.iter_cnt, self.acc_reward))

        return observation, reward, done, {"submit": submit, "return": self.acc_reward}

    def reset(self):
        self.submitted.clear()

        self.iter_cnt = 0
        self.acc_reward = 0

        self.board = np.copy(self.board_backup)

        observation = np.copy(self.board)  # .get()
        return observation

    def render(self, mode='ansi', close=False):
        return ""

    def seed(self, seed=None):
        # with cp.cuda.Device(self.cupy_device):
        np.random.seed(seed)

    def close(self):
        # if self.play:
        #     self.play_log.close()
        return


from pyinrail import pyinrail
import pytesseract
import datetime
import pandas as pd

pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

# Get the current date
current_date = datetime.date.today()

# Calculate the ARP date by adding a certain number of days to the current date
arp_date = current_date + datetime.timedelta(days=0)  # Adjust the number of days as needed
print(arp_date)

# Convert the ARP date to the required format
arp_date_str = arp_date.strftime("%d-%m-%Y")

# Create a list of stations
stations = ['KOLLAM JN', 'TAMBARAM', 'CHENNAI EGMORE']  # Add the additional station here

# Perform a railway enquiry
enq = pyinrail.RailwayEnquiry(src=stations[0], dest=stations[-1], date=arp_date_str)
df = enq.get_trains_between_stations(as_df=True)
station_name = enq.stations
# print(station_name)
print(df)

# # Filter the DataFrame to get trains passing through Alappuzha
# list_trains = df[df['Stations'].str.contains('Alappuzha', case=False)]
#
# # Convert 'Departure_Time' column to datetime.time objects
# list_trains['Departure_Time'] = pd.to_datetime(list_trains['Departure_Time'], format='%H:%M').dt.time
#
# # Find the train with departure time closest to 11 am
# target_time = datetime.time(11, 0)  # Desired time: 11 am
# closest_train = list_trains.iloc[(list_trains['Departure_Time'] - target_time).abs().argsort()[0]]
#
# # Print the details of the train
# print("Train Details:")
# print("Train Number:", closest_train['Train_Number'])
# print("Train Name:", closest_train['Train_Name'])
# print("Departure Time:", closest_train['Departure_Time'])


import pandas as pd
import glob
import os
import numpy as np
from pathlib import Path

# get data file names
software_path = r'data/softwares' # use your path
software_files = glob.glob(software_path + "/*")

#----------Check result_base.csv exist -----------
outPutPath = Path('data/base_data/result_base.csv')
software_old_base_dict = pd.read_csv(outPutPath).set_index('software.name',drop=False).to_dict("index") if outPutPath.exists() else {}

#------------End -----------------------

#----------Check need update------------------
updatePath = r'data/update'
update_files = glob.glob(updatePath + "/*")
update_df_list = []
for update_file in update_files:
    try:
        df = pd.read_csv(update_file)
        update_df_list.append(df)
    except:
        df = pd.read_excel(update_file)
        update_df_list.append(df)

#------------End -----------------------

software_df_list = []
software_name_dict = {}

for software_file in software_files:
    print("Read:" + software_file + '\n')
    sub_cat, main_cat= "",""
    try:
        df = pd.read_csv(software_file)
    except:
        df = pd.read_excel(software_file)
    df = df.dropna(subset=['software.name'])
    df = df.drop_duplicates(subset=['software.name'])
    software_df_list.append(df)
    dict = df.to_dict('records')
    try:
        sub_cat, main_cat = software_file.split('/')[-1].split(', ')
        main_cat = main_cat.split('.')[0]
        for row in dict:
            try:
                software_name_dict[str(row['software.name']).rstrip()] = [sub_cat,main_cat]
            except:
                continue
    except:
        for row in dict:
            try:
                software_name_dict[str(row['software.name']).rstrip()] = [str(row['sub.category']),str(row['main.category'])]
            except:
                continue
software_new_df = pd.DataFrame()
for df in software_df_list:
    software_new_df = pd.concat([software_new_df,df], axis = 0, ignore_index = True)

software_new_df = pd.DataFrame(software_new_df,columns = ['software.name','logo','Company','Description','Website','main.category','Sub.cat1','Sub.cat2','Pricing','entry.price','liked','match.score','x2.rating','overall.rating','total.reviews','recomm.rating','recomm.count','easeofuse.rating','easeofuse.count','cust.serv.rating','cust.serv.count','feature.rating','feature.count','valuemoney.rating','valuemoney.count','pros.count','cons.count','overall.1rating','overall.2rating','overall.3rating','overall.4rating','overall.5rating',
'free.trial','subscription','free.version','free.demo','open.source','train.inperson','train.online','train.webinar','train.liverep','train.doc','support.hours','support.online','support.liverep','feature_list','location','industry','sector','employee','tel','email','contact','demo','freetrial','getstarted','download','document','tutorial','community','feature','github','claim','support.url','pricing.url'])
software_new_df = software_new_df.drop_duplicates(subset=['software.name']).set_index("software.name",drop=False)

unmodified_columns = {'software.name','logo','Company','Description','Website','main.category','Sub.cat1','Sub.cat2','Pricing','entry.price','subscription','free.version','free.demo','open.source','train.inperson','train.online','train.webinar','train.liverep','train.doc','support.hours','support.online','support.liverep','feature_list','location','industry','sector','employee','tel','email','contact','demo','freetrial','getstarted','download','document','tutorial','community','feature','github'}
software_new_dict = software_new_df.to_dict('records')
for row in software_new_dict:
    row['software.name'] = str(row['software.name']).rstrip()
    if str(row['Company']).find('By ')!=-1:
        row['Company'] = str(row['Company']).replace('By ','')
    if str(row['software.name']) in software_name_dict.keys():
        row['Sub.cat1'] = software_name_dict[str(row['software.name'])][0]
        row['main.category'] = software_name_dict[str(row['software.name'])][1]

    if row['software.name'] not in software_old_base_dict.keys():
        software_old_base_dict[row['software.name']] = row
    else:
        for col,val in row.items():
            if col in unmodified_columns and str(software_old_base_dict[row['software.name']][col])!='nan':
                continue
            if str(val)!='nan':
                software_old_base_dict[row['software.name']][col] = val

#--------Check update folder files -------------------------
for update_df in update_df_list:
    update_df_dict = update_df.set_index('software.name',drop=False).to_dict("records") if 'software.name' in set(update_df.columns) else {}
    for row in update_df_dict:
        row['software.name'] = str(row['software.name']).rstrip()
        if str(row['software.name']) not in software_old_base_dict.keys():
            software_old_base_dict[str(row['software.name'])] = {}
        for col,val in row.items():
            if str(col) in set(software_new_df.columns):
                if str(row['software.name']) in software_old_base_dict:
                    software_old_base_dict[row['software.name']][col] = val
#------------End-------------------------------------
software_base = pd.DataFrame.from_dict(software_old_base_dict,orient='index')
software_base.replace(r'^\s*$', np.nan, regex=True)
outPutPath = Path('data/base_data/result_base.csv')
software_base.to_csv(outPutPath,index=False)


'''
This is the graveyard for old Sciris functions that are no longer used. It is preserved
for backwards compatibility, but is not imported into Sciris by default.
'''

import types
import warnings
import textwrap
import traceback
import gzip as gz
import numpy as np
import pickle as pkl
import datetime as dt
import copyreg as cpreg
import multiprocess as mp
from io import BytesIO as IO
from functools import partial
from contextlib import closing
from .. import sc_utils as scu
from .. import sc_odict as sco
from .. import sc_fileio as scf
from .. import sc_profiling as scp

##############################################################################
#%% Python 2 legacy support
##############################################################################

not_string_pickleable = ['datetime', 'BytesIO']
byte_objects = ['datetime', 'BytesIO', 'odict', 'spreadsheet', 'blobject']

def loadobj2or3(filename=None, filestring=None, recursionlimit=None, **kwargs):  # pragma: no cover
    '''
    Try to load as a (Sciris-saved) Python 3 pickle; if that fails, try to load
    as a Python 2 pickle. For legacy support only.

    For available keyword arguments, see sc.load().

    Args:
        filename (str): the name of the file to load
        filestring (str): alternatively, specify an already-loaded bytestring
        recursionlimit (int): how deeply to parse objects before failing (default 1000)
    '''
    try:
        output = scf.loadobj(filename=filename, **kwargs)
    except:
        output = _loadobj2to3(filename=filename, filestring=filestring, recursionlimit=recursionlimit)
    return output

def _loadobj2to3(filename=None, filestring=None, recursionlimit=None): # pragma: no cover
    '''
    Used by loadobj2or3() to load Python2 objects in Python3 if all other
    loading methods fail. Uses a recursive approach, so can set a recursion limit.
    '''

    class Placeholder():
        ''' Replace these corrupted classes with properly loaded ones '''
        def __init__(*args):
            return

        def __setstate__(self, state):
            if isinstance(state,dict):
                self.__dict__ = state
            else:
                self.state = state
            return

    class StringUnpickler(pkl.Unpickler):
        def find_class(self, module, name, verbose=False):
            if verbose: print('Unpickling string module %s , name %s' % (module, name))
            if name in not_string_pickleable:
                return scf.Empty
            else:
                try:
                    output = pkl.Unpickler.find_class(self,module,name)
                except Exception as E:
                    print('Warning, string unpickling could not find module %s, name %s: %s' % (module, name, str(E)))
                    output = scf.Empty
                return output

    class BytesUnpickler(pkl.Unpickler):
        def find_class(self, module, name, verbose=False):
            if verbose: print('Unpickling bytes module %s , name %s' % (module, name))
            if name in byte_objects:
                try:
                    output = pkl.Unpickler.find_class(self,module,name)
                except Exception as E:
                    print('Warning, bytes unpickling could not find module %s, name %s: %s' % (module, name, str(E)))
                    output = Placeholder
                return output
            else:
                return Placeholder

    def recursive_substitute(obj1, obj2, track=None, recursionlevel=0, recursionlimit=None):
        if recursionlimit is None: # Recursion limit
            recursionlimit = 1000 # Better to die here than hit Python's recursion limit

        def recursion_warning(count, obj1, obj2):
            output = 'Warning, internal recursion depth exceeded, aborting: depth=%s, %s -> %s' % (count, type(obj1), type(obj2))
            return output

        recursionlevel += 1

        if track is None:
            track = []

        if isinstance(obj1, scf.Blobject): # Handle blobjects (usually spreadsheets)
            obj1.blob  = obj2.__dict__[b'blob']
            obj1.bytes = obj2.__dict__[b'bytes']

        if isinstance(obj2, dict): # Handle dictionaries
            for k,v in obj2.items():
                if isinstance(v, dt.datetime):
                    setattr(obj1, k.decode('latin1'), v)
                elif isinstance(v, dict) or hasattr(v,'__dict__'):
                    if isinstance(k, (bytes, bytearray)):
                        k = k.decode('latin1')
                    track2 = track.copy()
                    track2.append(k)
                    if recursionlevel<=recursionlimit:
                        recursionlevel = recursive_substitute(obj1[k], v, track2, recursionlevel, recursionlimit)
                    else:
                        print(recursion_warning(recursionlevel, obj1, obj2))
        else:
            for k,v in obj2.__dict__.items():
                if isinstance(v, dt.datetime):
                    setattr(obj1,k.decode('latin1'), v)
                elif isinstance(v,dict) or hasattr(v,'__dict__'):
                    if isinstance(k, (bytes, bytearray)):
                        k = k.decode('latin1')
                    track2 = track.copy()
                    track2.append(k)
                    if recursionlevel<=recursionlimit:
                        recursionlevel = recursive_substitute(getattr(obj1,k), v, track2, recursionlevel, recursionlimit)
                    else:
                        print(recursion_warning(recursionlevel, obj1, obj2))
        return recursionlevel

    def loadintostring(fileobj):
        unpickler1 = StringUnpickler(fileobj, encoding='latin1')
        try:
            stringout = unpickler1.load()
        except Exception as E:
            print('Warning, string pickle loading failed: %s' % str(E))
            exception = traceback.format_exc() # Grab the trackback stack
            stringout = scf.makefailed(module_name='String unpickler failed', name='n/a', error=E, exception=exception)
        return stringout

    def loadintobytes(fileobj):
        unpickler2 = BytesUnpickler(fileobj,  encoding='bytes')
        try:
            bytesout  = unpickler2.load()
        except Exception as E:
            print('Warning, bytes pickle loading failed: %s' % str(E))
            exception = traceback.format_exc() # Grab the trackback stack
            bytesout = scf.makefailed(module_name='Bytes unpickler failed', name='n/a', error=E, exception=exception)
        return bytesout

    # Load either from file or from string
    if filename:
        with gz.GzipFile(filename) as fileobj:
            stringout = loadintostring(fileobj)
        with gz.GzipFile(filename) as fileobj:
            bytesout = loadintobytes(fileobj)

    elif filestring:
        with closing(IO(filestring)) as output:
            with gz.GzipFile(fileobj=output, mode='rb') as fileobj:
                stringout = loadintostring(fileobj)
        with closing(IO(filestring)) as output:
            with gz.GzipFile(fileobj=output, mode='rb') as fileobj:
                bytesout = loadintobytes(fileobj)
    else:
        errormsg = 'You must supply either a filename or a filestring for loadobj() or loadstr(), respectively'
        raise Exception(errormsg)

    # Actually do the load, with correct substitution
    recursive_substitute(stringout, bytesout, recursionlevel=0, recursionlimit=recursionlimit)
    return stringout

##############################################################################
#%% Twisted pickling methods
##############################################################################

# NOTE: The code below is part of the Twisted package, and is included
# here to allow functools.partial() objects (among other things) to be
# pickled; they are not for public consumption. --CK

# From: twisted/persisted/styles.py
# -*- test-case-name: twisted.test.test_persisted -*-
# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

_UniversalPicklingError = pkl.PicklingError

def _pickleMethod(method):
    return (_unpickleMethod, (method.__name__,         method.__self__, method.__self__.__class__))

def _methodFunction(classObject, methodName):
    methodObject = getattr(classObject, methodName)
    return methodObject

def _unpickleMethod(im_name, im_self, im_class):
    if im_self is None:
        return getattr(im_class, im_name)
    try:
        methodFunction = _methodFunction(im_class, im_name)
    except AttributeError: # pragma: no cover
        assert im_self is not None, "No recourse: no instance to guess from."
        if im_self.__class__ is im_class:
            raise
        return _unpickleMethod(im_name, im_self, im_self.__class__)
    else:
        maybeClass = ()
        bound = types.MethodType(methodFunction, im_self, *maybeClass)
        return bound

cpreg.pickle(types.MethodType, _pickleMethod, _unpickleMethod)

# Legacy support for loading Sciris <1.0 objects; may be removed in future
pickleMethod = _pickleMethod
unpickleMethod = _unpickleMethod

##############################################################################
#%% Legacy data frame class
##############################################################################

class legacy_dataframe(object): # pragma: no cover
    '''
    This legacy dataframe is maintained solely to allow loading old files.

    **Example**::

        import sciris as sc
        from sciris.extras import legacy as scl
        remapping = {'sciris.sc_dataframe.dataframe':scl.legacy_dataframe}
        old = sc.load('my-old-file.obj', remapping=remapping)

    | Version: 2020nov29
    | Migrated to ``sc_legacy`` in version 2.0.0.
    '''

    def __init__(self, cols=None, data=None, nrows=None):
        self.cols = None
        self.data = None
        self.make(cols=cols, data=data, nrows=nrows)
        return

    def __repr__(self, spacing=2):
        ''' spacing = space between columns '''
        if not self.cols: # No keys, give up
            return '<empty dataframe>'

        else: # Go for it
            outputlist = sco.odict()
            outputformats = sco.odict()

            # Gather data
            nrows = self.nrows
            for c,col in enumerate(self.cols):
                outputlist[col] = list()
                maxlen = len(col) # Start with length of column name
                if nrows:
                    for val in self.data[:,c]:
                        output = scu.flexstr(val)
                        maxlen = max(maxlen, len(output))
                        outputlist[col].append(output)
                outputformats[col] = '%'+'%i'%(maxlen+spacing)+'s'

            ndigits = (np.floor(np.log10(max(1,nrows)))+1) # Don't allow 0 rows
            indformat = '%%%is' % ndigits # Choose the right number of digits to print

            # Assemble output
            output = indformat % '' # Empty column for index
            for col in self.cols: # Print out header
                output += outputformats[col] % col
            output += '\n'

            for ind in range(nrows): # Loop over rows to print out
                output += indformat % scu.flexstr(ind)
                for col in self.cols: # Print out data
                    output += outputformats[col] % outputlist[col][ind]
                if ind<nrows-1: output += '\n'

            return output

    @property
    def ncols(self):
        ''' Get the number of columns in the data frame '''
        ncols = len(self.cols)
        ncols2 = self.data.shape[1]
        if ncols != ncols2:
            errormsg = 'Dataframe corrupted: %s columns specified but %s in data' % (ncols, ncols2)
            raise Exception(errormsg)
        return ncols

    @property
    def nrows(self):
        ''' Get the number of rows in the data frame '''
        try:    return self.data.shape[0]
        except: return 0 # If it didn't work, probably because it's empty

    @property
    def shape(self):
        ''' Equivalent to the shape of the data array, minus the headers '''
        return (self.nrows, self.ncols)

    def make(self, cols=None, data=None, nrows=None):
        '''
        Creates a dataframe from the supplied input data.

        **Usage examples**::

            df = sc.dataframe()
            df = sc.dataframe(['a','b','c'])
            df = sc.dataframe(['a','b','c'], nrows=2)
            df = sc.dataframe([['a','b','c'],[1,2,3],[4,5,6]])
            df = sc.dataframe(['a','b','c'], [[1,2,3],[4,5,6]])
            df = sc.dataframe(cols=['a','b','c'], data=[[1,2,3],[4,5,6]])
        '''
        import pandas as pd # Optional import

        # Handle columns
        if nrows is None:
            nrows = 0
        if cols is None and data is None:
            cols = list()
            data = np.zeros((int(nrows), 0), dtype=object) # Object allows more than just numbers to be stored
        elif cols is None and data is not None: # Shouldn't happen, but if it does, swap inputs
            cols = data
            data = None

        if isinstance(cols, pd.DataFrame): # It's actually a Pandas dataframe
            self.pandas(df=cols)
            return # We're done

        # A dictionary is supplied: assume keys are columns, and the rest is the data
        if isinstance(cols, dict):
            data = [col for col in cols.values()]
            cols = list(cols.keys())

        elif not scu.checktype(cols, 'listlike'):
            errormsg = 'Inputs to dataframe must be list, tuple, or array, not %s' % (type(cols))
            raise Exception(errormsg)

        # Handle data
        if data is None:
            if np.ndim(cols)==2 and np.shape(cols)[0]>1: # It's a 2D array with more than one row: treat first as header
                data = scu.dcp(cols[1:])
                cols = scu.dcp(cols[0])
            else:
                data = np.zeros((int(nrows),len(cols)), dtype=object) # Just use default
        data = np.array(data, dtype=object)
        if data.ndim != 2:
            if data.ndim == 1:
                if len(cols)==1: # A single column, use the data to populate the rows
                    data = np.reshape(data, (len(data),1))
                elif len(data)==len(cols): # A single row, use the data to populate the columns
                    data = np.reshape(data, (1,len(data)))
                else:
                    errormsg = 'Dimension of data can only be 1 if there is 1 column, not %s' % len(cols)
                    raise Exception(errormsg)
            else:
                errormsg = 'Dimension of data must be 1 or 2, not %s' % data.ndim
                raise Exception(errormsg)
        if data.shape[1]==len(cols):
            pass
        elif data.shape[0]==len(cols):
            data = data.transpose()
        else:
            errormsg = 'Number of columns (%s) does not match array shape (%s)' % (len(cols), data.shape)
            raise Exception(errormsg)

        # Store it
        self.cols = list(cols)
        self.data = data
        return

##############################################################################
#%% Parallelization
##############################################################################

def parallelcmd(cmd=None, parfor=None, returnval=None, maxcpu=None, maxmem=None, interval=None, die=True, **kwargs):
    '''
    A function to parallelize any block of code. Note: this is intended for quick
    prototyping only; since it uses exec(), it is not recommended for use in production
    code.

    Args:
        cmd       (str):   a string representation of the code to be run in parallel
        parfor    (dict):  a dictionary of lists of the variables to loop over
        returnval (str):   the name of the output variable
        maxcpu    (float): maximum CPU load; used by ``sc.loadbalancer()``
        maxmem    (float): maximum fraction of virtual memory (RAM); used by ``sc.loadbalancer()``
        interval  (float): the time delay to poll to see if load is OK,  used in ``sc.loadbalancer()``
        die       (bool):  whether to stop immediately if an exception is encountered (otherwise, store the exception as the result)
        kwargs    (dict):  variables to pass into the code

    **Example**::

        const = 4
        parfor = {'val':[3,5,9]}
        returnval = 'result'
        cmd = """
        newval = val+const
        result = newval**2
        """
        results = sc.parallelcmd(cmd=cmd, parfor=parfor, returnval=returnval, const=const)

    | New in version 2.0.0: replaced ``maxload`` with ``maxcpu``/``maxmem``; automatically de-indent the command
    | Migrated to ``sc_legacy`` in version 2.2.0.
    '''

    # Handle maxload
    maxload = kwargs.pop('maxload', None)
    if maxload is not None: # pragma: no cover
        maxcpu = maxload
        warnmsg = 'sc.loadbalancer() argument "maxload" has been renamed "maxcpu" as of v2.0.0'
        warnings.warn(warnmsg, category=FutureWarning, stacklevel=2)

    # Deindent the command
    cmd = textwrap.dedent(cmd)

    # Create queue
    nfor = len(list(parfor.values())[0])
    outputqueue = mp.Queue()
    outputlist = np.empty(nfor, dtype=object)
    processes = []
    for i in range(nfor):
        args = (cmd, parfor, returnval, i, outputqueue, maxcpu, maxmem, interval, die, kwargs)
        prc = mp.Process(target=_parallelcmd_task, args=args)
        prc.start()
        processes.append(prc)
    for i in range(nfor):
        _i,returnval = outputqueue.get()
        outputlist[_i] = returnval
    for prc in processes:
        prc.join() # Wait for them to finish

    outputlist = outputlist.tolist()

    return outputlist

def parallel_progress(fcn, inputs, num_workers=None, show_progress=True, initializer=None): # pragma: no cover
    """
    Run a function in parallel with a optional single progress bar

    The result is essentially equivalent to::

        >>> list(map(fcn, inputs))

    But with execution in parallel and with a single progress bar being shown.

    Args:
        fcn (function): Function object to call, accepting one argument, OR a function with zero arguments in which case inputs should be an integer
        inputs (list): A collection of inputs that will each be passed to the function OR a number, if the fcn() has no input arguments
        num_workers (int): Number of processes, defaults to the number of CPUs
        show_progress (bool): Whether to show a progress bar
        initializer (func): A function that each worker process will call when it starts

    Returns:
        A list of outputs

    | New in version 1.0.0.
    | Migrated to ``sc_legacy`` in version 2.2.0.
    """
    try:
        from tqdm import tqdm
    except ModuleNotFoundError as E:
        errormsg = 'Module tqdm not found; please install with "pip install tqdm"'
        raise ModuleNotFoundError(errormsg) from E

    pool = mp.Pool(num_workers, initializer=initializer)

    results = [None]
    if scu.isnumber(inputs):
        results *= inputs
        pbar = tqdm(total=inputs) if show_progress else None
    else:
        results *= len(inputs)
        pbar = tqdm(total=len(inputs)) if show_progress else None

    def callback(result, idx):
        results[idx] = result
        if show_progress:
            pbar.update(1)

    if scu.isnumber(inputs):
        for i in range(inputs):
            pool.apply_async(fcn, callback=partial(callback, idx=i))
    else:
        for i, x in enumerate(inputs):
            pool.apply_async(fcn, args=(x,), callback=partial(callback, idx=i))

    pool.close()
    pool.join()

    if show_progress:
        pbar.close()

    return results

def _parallelcmd_task(_cmd, _parfor, _returnval, _i, _outputqueue, _maxcpu, _maxmem, _interval, _die, _kwargs): # pragma: no cover # No coverage since pickled
    '''
    The task to be executed by ``sc.parallelcmd()``. All internal variables start with
    underscores to avoid possible collisions in the ``exec()`` statements. Not to be called
    directly.

    Migrated to ``sc_legacy`` in version 2.2.0.
    '''
    if _maxcpu or _maxmem:
        scp.loadbalancer(maxcpu=_maxcpu, maxmem=_maxmem, index=_i, interval=_interval)

    # Set the loop variables
    for _key in _parfor.keys():
        _thisval = _parfor[_key][_i] # analysis:ignore
        exec(f'{_key} = _thisval') # Set the value of this variable

    # Set the keyword arguments
    for _key in _kwargs.keys():
        _thisval = _kwargs[_key] # analysis:ignore
        exec(f'{_key} = _thisval') # Set the value of this variable

    # Run the command
    try:
        exec(_cmd) # The meat of the matter!
    except Exception:
        if _die:
            raise Exception
        else:
            warnmsg = f'sc.parallelcmd(): Task {_i} failed, but die=False so continuing.\n{scu.traceback()}'
            warnings.warn(warnmsg, category=RuntimeWarning, stacklevel=2)
            exec(f'{_returnval} = None')

    # Append results
    _outputqueue.put((_i,eval(_returnval)))

    return


age = input("How old are you? ")
if age:
    age = int(age)
    if int(age) >= 18 and int(age) <21:
        print("You either, but need a wristband!")
        # 18-21 wristband
    elif int(age) >= 21:
        print("You are good to enter and can drink!")
        #21+ drink, normal entry
    else:
        # too young, sorry
        print("You can't come in little one")

else:
    print("please enter an age")


# SQL Models
from typing import List

from aiosqlite import Connection

class DBModel:
    """
    ABC To allow for unified use of database models.
    """

    _state: Connection
    __rows__: list
    _data: dict

    @classmethod
    async def create_table(cls, connection: Connection, *rows, name=None) -> None:
        """Creates the table.

        :param connection: The aiosqlite connection object.
        :param rows: A list of rows to create (name type constraints)
        :param name: Optional[str] - The name of the table. Defaults to cls.__class__.__name__."""
        if not rows:
            rows = cls.__rows__
        query = "CREATE TABLE IF NOT EXISTS {} ({});".format(name or cls.__class__.__name__.lower(), ",\n".join(rows))
        # The IF NOT EXISTS in there means this function can be called numerous times.
        await connection.execute(query)
        await connection.commit()

    @classmethod
    async def create(cls, state, **values):
        """
        [C]RUD - Creates an entry in the database.

        :param values: column:value pairs to use.
        :param state: The state to use
        :return: Resolved model
        """
        raise NotImplementedError

    @classmethod
    async def get(cls, state, key, **other_comps):
        """
        C[R]UD - Retrieve an entry from the database.

        :param state: The state to use
        :param key: The primary key to choose from, if **other_comps is not provided.
        :param other_comps: A column:value pair.
        :return: Resolved model
        """
        raise NotImplementedError

    async def edit(self, keys: List[str], values) -> None:
        """
        CR[U]D - Edits the entry on the database.

        .. warning::
            This function commits directly to the database automatically. Make sure your datatypes are correct.

        :param keys: The column names to update
        :param values: the values to update with
        :return: None
        """
        raise NotImplementedError

    async def delete(self) -> None:
        """
        CRU[D] - Deletes this entry.

        .. warning::
            This is irreversible!

        :return: Nothing. You won't be able to use this model again after.
        """
        raise NotImplementedError

    def __getattr__(self, item):
        if not isinstance(item, str):
            raise TypeError("Got type '{.__class__.__name__}' instead of 'str'.".format(item))
        try:
            return self._data[item]
        except KeyError:
            raise AttributeError(self.__class__.__name__ + " has no attribute '{}'.".format(item))

    __getitem__ = __getattr__  # We'll just alias getitem to getattr (so we can do stuff like model["key"] or model.key)

    def __setattr__(self, key, value):
        # We can only edit self, setting attributes will just do it to the local object.
        # This is bad because we use an async lib, so can't do the update here ourselves.
        # And we also don't want to end up with a race condition where we have tasks competing to see who can
        # edit the database the fastest. Asyncio tasks really do just show off in that regard.
        raise AttributeError(f"Please use {self.__class__.__name__}.edit() over setting attributes.")

    @property
    def state(self) -> Connection:
        """
        The connection state

        :return: aiosqlite.Connection - the current database connection
        """
        return self._state

    @property
    def data(self) -> dict:
        """
        The raw dictionary of data for this table.

        :return: dict - key: value pairs
        """
        return self._data.copy()

class Guild(DBModel):
    """
    A database model representing a guild.
    """

    __rows__ = [
        "id INTEGER PRIMARY KEY NOT NULL UNIQUE",
        "prefix TEXT NULLABLE DEFAULT NULL",
        "mod_log INTEGER UNIQUE",
        "mod_role INTEGER UNIQUE",
        "case_id INTEGER",
    ]

    def __init__(self, data: dict, *, state):
        self._state = state
        self._data = data

    @classmethod
    async def create(cls, state: Connection, **kwargs):
        query = """
        INSERT INTO guilds VALUES (?, ?, ?, ?, ?);
        """
        await state.execute(query, *kwargs.values())
        await state.commit()
        return cls(kwargs, state=state)

    @classmethod
    async def get(cls, state, key, **other_comps):
        query = """
        SELECT * FROM guilds WHERE {} LIMIT 1;
        """
        if other_comps:
            where = " AND ".join(f"{column}={value}" for column, value in other_comps.items())
        else:
            where = "id=" + key
        query = query.format(where)
        data = {}
        async with state.execute(query) as cursor:
            row = await cursor.fetchone()
            for key in row.keys():
                data[key] = row[key]
        return cls(data, state=state)

    async def edit(self, keys: List[str], values):
        if not hasattr(self, "id"):
            raise ValueError("Missing ID for guild data model - unable to update")
        query = """
        UPDATE guilds
        SET {}
        WHERE id=?"""

        _set = []
        for n, key in enumerate(keys):
            _set.append(f"{key}={values[n]}")
        query = query.format(", ")
        await self.state.execute(query, (self.id,))
        await self.state.commit()

    async def delete(self) -> None:
        if not hasattr(self, "id"):
            raise ValueError("Missing ID for guild data model - unable to delete.")
        await self.state.execute(
            """
            DELETE FROM guilds
            WHERE id=?
            """,
            (self.id,),
        )
        await self.state.commit()
        return