import datetime
import json

from django.core.exceptions import ValidationError
from django.db import IntegrityError
from django.db.models import Count
from django.http import JsonResponse
from django.utils.decorators import method_decorator
from django.views.decorators.csrf import csrf_exempt
from rest_framework.authentication import TokenAuthentication
from rest_framework.permissions import IsAuthenticated
from rest_framework.views import APIView

from restauth.models import CustomUser
from .models import Place, Placedoes, JobRequest, Job, Service, DaySchedule, Provides, Appointment, Serviceinstance, \
    Promotion, Promoincludes, Jobtype, ScheduleTemplate
from .permissions import IsOwner, IsProvider, IsProviderPro
from .serializers import PlaceSerializer, JobRequestSerializer, ServiceSerializer, JobSerializer, \
    ServiceInstanceSerializer, AppointmentSerializer, DayScheduleSerializer, ProvidesSerializer, \
    ScheduleTemplateSerializer

# custom api view implementing get_object with permission check
class CustomAPIView(APIView):
    authentication_classes = [TokenAuthentication]
    responseOK = {'result': 'OK'}

    def get_object(self, request, obj_model, pk):
        obj = obj_model.objects.get(pk=pk)
        self.check_object_permissions(request, obj)
        return obj

    def returnOK(self):
        return JsonResponse(self.responseOK)

    def returnError(self, code, message):
        return JsonResponse({'reason': message}, status=code)

@method_decorator(csrf_exempt, name='dispatch')
class UserProfile(APIView):
    permission_classes = [IsAuthenticated & IsProvider]

    def get(self, request):
        return JsonResponse({
            'isprovider': request.user.isprovider,
            'isclient': request.user.isclient,
        })

@method_decorator(csrf_exempt, name='dispatch')
class NewPlace(CustomAPIView):
    permission_classes = [IsAuthenticated & IsProvider]

    # street, streetnumber, apnumber, city, state, country, businessname, phonenumber, email
    def post(self, request):
        params = request.POST

        place = Place()
        place.serviceprovider = request.user
        place.address = params.get('address')
        place.businessname = params.get('businessname')
        place.phonenumber = params.get('phonenumber')
        place.email = ""

        place.save()

        if params.get('works_here') == "true":
            job = Job()
            job.serviceprovider = request.user
            job.place = place
            job.save()

        return JsonResponse({'place_id': place.id})

@method_decorator(csrf_exempt, name='dispatch')
class DropPlace(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    # place_id
    def post(self, request):
        place = self.get_object(request, Place, request.POST.get('place_id'))
        place.delete()

        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class JobTypesView(CustomAPIView):
    # permission_classes = [IsAuthenticated & IsOwner]

    def get(self, request):
        jobtypes_list = Jobtype.objects.all()

        name_list = []
        for x in jobtypes_list:
            name_list.append(x.type)

        return JsonResponse({'jobtypes': name_list})

@method_decorator(csrf_exempt, name='dispatch')
class PlaceDoes(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    # place, jobtype
    def post(self, request):
        body = json.loads(request.body)
        place = self.get_object(request, Place, body['place'])
        jobtypes = body['jobtypes']

        placedoes_created = []
        for x in jobtypes:
            placedoes = Placedoes()
            jobtype = Jobtype.objects.get(type=x)
            placedoes.jobtype = jobtype
            placedoes.place = place

            try:
                placedoes.full_clean()
                placedoes_created.append(placedoes)
            except IntegrityError:
                return self.returnError(400, x)

        for x in placedoes_created:
            x.save()

        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class PlacesOwned(CustomAPIView):
    permission_classes = [IsAuthenticated]

    def get(self, request):
        places = Place.objects.filter(serviceprovider=request.user.id)
        output = []

        for p in places:
            output.append(PlaceSerializer(p).data)
        return JsonResponse({'places': output})

@method_decorator(csrf_exempt, name='dispatch')
class JobsView(CustomAPIView):
    permission_classes = [IsAuthenticated & IsProvider]

    def get(self, request):
        jobs = Job.objects.filter(serviceprovider=request.user)
        output = {}
        for j in jobs:
            schedules_list = {}
            day_schedules = DaySchedule.objects.filter(job=j)
            for ds in day_schedules:
                provides_list = []
                provides = Provides.objects.filter(job=j, day_schedule=ds)

                for p in provides:
                    provides_list.append(ProvidesSerializer(p).data)

                schedules_list[ds.day_of_week] = {
                    "day_schedule": DayScheduleSerializer(ds).data,
                    "provides": provides_list
                }

            output[j.id] = {
                "job": JobSerializer(j).data,
                "day_schedules": schedules_list
            }

        return JsonResponse({'jobs': output})

@method_decorator(csrf_exempt, name='dispatch')
class SearchPlace(CustomAPIView):
    permission_classes = [IsAuthenticated]

    def apply_transformation(self, string):
        return string.strip().lower()

    def get(self, request):
        businessname = request.GET.get('searchquery')
        businessname = self.apply_transformation(businessname)
        output = []
        print(businessname)
        if businessname is not None and businessname != '':
            places = Place.objects.filter(businessname__icontains=businessname)
            for p in places:
                if p.serviceprovider != request.user:
                    output.append(PlaceSerializer(p).data)
        return JsonResponse({'places': output})

@method_decorator(csrf_exempt, name='dispatch')
class NewJobRequest(CustomAPIView):
    permission_classes = [IsAuthenticated & IsProvider]

    def post(self, request):
        place_id = request.POST.get('place_id')

        if Job.objects.filter(place=place_id, serviceprovider=request.user.id).exists():
            return self.returnError(500, 'Job already exists')

        place = Place.objects.get(id=place_id)

        job_request = JobRequest()
        job_request.place = place
        job_request.serviceprovider = request.user
        job_request.save()
        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class ViewJobRequests(CustomAPIView):
    permission_classes = [IsAuthenticated]

    def get(self, request):
        places = Place.objects.filter(serviceprovider=request.user.id)
        output = {}
        for p in places:
            job_requests = JobRequest.objects.filter(place=p.id)
            jb_list = []
            for jr in job_requests:
                jb_list.append(JobRequestSerializer(jr).data)
            output[p.id] = jb_list

        return JsonResponse(output)

@method_decorator(csrf_exempt, name='dispatch')
class AcceptJobRequest(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        place_id = request.POST.get('place_id')
        from_who = request.POST.get('serviceprovider')

        user_from = CustomUser.objects.get(email=from_who)
        place = self.get_object(request, Place, place_id)

        try:
            jr = JobRequest.objects.get(place=place, serviceprovider=user_from)

            job = Job()
            job.serviceprovider = user_from
            job.place = place

            job.save()
            jr.delete()

            return self.returnOK()
        except JobRequest.DoesNotExist:
            raise

@method_decorator(csrf_exempt, name='dispatch')
class CancelJobRequest(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        place_id = request.POST.get('place_id')
        from_who = request.POST.get('serviceprovider')

        user_from = CustomUser.objects.get(email=from_who)
        place = self.get_object(request, Place, place_id)

        try:
            jr = JobRequest.objects.get(place=place, serviceprovider=user_from)
            jr.delete()

            return self.returnOK()
        except JobRequest.DoesNotExist:
            raise

@method_decorator(csrf_exempt, name='dispatch')
class DoableServices(CustomAPIView):
    permission_classes = [IsAuthenticated & IsProvider]

    def get(self, request):
        place_id = request.GET.get('place_id')

        output = {}
        jobtypes = Placedoes.objects.filter(place_id=place_id)
        for pd in jobtypes:

            services = Service.objects.filter(jobtype=pd.jobtype.type)
            services_json = []
            for s in services:
                services_json.append(ServiceSerializer(s).data)

            output[pd.jobtype.type] = services_json

        return JsonResponse(output)

@method_decorator(csrf_exempt, name='dispatch')
class NewDaySchedule(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        body = json.loads(request.body)
        job = self.get_object(request, Job, body['job_id'])

        job_day_schedules = DaySchedule.objects.filter(job=job)
        for ds in job_day_schedules:
            if str(ds.day_of_week) not in body['days']:
                ds.delete()

        for day in body['days']:
            day_schedule = DaySchedule.objects.filter(job=job, day_of_week=day)
            config = body['days'][day]

            pause_start = config['pause_start'] if 'pause_start' in config else None
            pause_end = config['pause_end'] if 'pause_end' in config else None

            if len(day_schedule) == 0:
                day_schedule = DaySchedule(
                    job=job,
                    day_of_week=int(day),
                    day_start=config['day_start'],
                    day_end=config['day_end'],
                    pause_start=pause_start,
                    pause_end=pause_end,
                )
            else:
                day_schedule = day_schedule[0]
                day_schedule.day_start = config['day_start']
                day_schedule.day_end = config['day_end']
                day_schedule.pause_start = pause_start
                day_schedule.pause_end = pause_end

            day_schedule.save()

        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class NewProvides(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        body = json.loads(request.body)
        job = self.get_object(request, Job, body['job_id'])

        day_schedule = DaySchedule.objects.get(job=job, day_of_week=body['day_of_week'])
        provides_in_dayschedule = Provides.objects.filter(day_schedule=day_schedule)

        for provides in provides_in_dayschedule:
            if str(provides.id) not in body['services']:
                provides.delete()

        provides_list = []
        for service_id in body['services']:
            p = Provides.objects.filter(day_schedule=day_schedule, service_id=service_id)

            if len(p) == 0:
                p = Provides(
                    job=job,
                    day_schedule=day_schedule,
                    service=Service.objects.get(id=int(service_id)),
                    cost=body['costs'][str(service_id)],
                    duration=body['durations'][str(service_id)],
                    parallelism=body['parallelisms'][str(service_id)],
                )
                p.full_clean()
            else:
                p = p[0]
                p.cost = body['costs'][str(service_id)]
                p.duration = body['durations'][str(service_id)],
                p.parallelism = body['parallelisms'][str(service_id)],

            provides_list.append(p)

        for x in provides_list:
            x.save()

        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class RemoveProvides(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        day_schedule = DaySchedule.objects.get(id=request.POST.get('day_schedule_id'))
        provides = Provides.objects.get(id=request.POST.get('provides_id'))
        self.get_object(request, Job, day_schedule.job.id)
        provides.remove()

        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class RemoveDaySchedule(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        day_schedule = DaySchedule.objects.get(id=request.POST.get('day_schedule_id'))
        self.get_object(request, Job, day_schedule.job.id)
        day_schedule.remove()
        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class DropJob(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        job = Job.objects.get(id=request.POST.get('job_id'))
        self.get_object(request, Place, job.place)
        job.delete()

@method_decorator(csrf_exempt, name='dispatch')
class GetAppointments(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def get(self, request):
        job_id = request.GET.get('job_id')
        job = Job.objects.get(id=job_id)

        if job.serviceprovider == request.user or job.place.serviceprovider == request.user:
            output = {}
            # TODO provide only future appointments
            appointments = Appointment.objects.filter(job=job, date__gte=datetime.date.today())
            for a in appointments:
                services = {}
                service_instances = Serviceinstance.objects.filter(appointment=a)
                for i, s in enumerate(service_instances):
                    services[i] = ServiceInstanceSerializer(s).data

                output[a.id] = {
                    'appointment': AppointmentSerializer(a).data,
                    'services': services
                }

            return JsonResponse(output)

        return self.returnError(403, "Not allowed to view appointments")

@method_decorator(csrf_exempt, name='dispatch')
class DropAppointment(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        appointment = Appointment.objects.get(request.POST.get('appointment_id'))
        job = self.get_object(request, Job, appointment.job)

        appointment.delete()

        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class NewPromotion(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        body = json.loads(request.body)
        job = Job.objects.get(id=body['job_id'])

        if job.serviceprovider == request.user or job.place.serviceprovider == request.user:
            since = body['since']
            to = body['to']
            description = body['description']
            weekday_list = body['weekday_list']
            services_ids = body['services']
            discounts = body['discounts']

            promo = Promotion(job=job, since=since, to=to, description=description, weekday_list=weekday_list)
            promo.save()
            includes_list = []
            for s in services_ids:
                service = Service.objects.get(id=s)
                promo_includes = Promoincludes(promotion=promo, service=service, discount=discounts[str(s)])
                try:
                    promo_includes.full_clean()
                    includes_list.append(promo_includes)
                except ValidationError:
                    promo.delete()
                    raise

            for x in includes_list:
                x.save()

            return self.returnOK()
        return self.returnError(403, 'Forbidden, you don\'t have permission to do that')

@method_decorator(csrf_exempt, name='dispatch')
class DropPromotion(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner]

    def post(self, request):
        promo = Promotion.objects.get(id=request.POST.body('promo_id'))
        self.get_object(request, Job, promo.job)
        promo.delete()
        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class GetFrequency(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner & IsProviderPro]

    def get(self, request):
        job = self.get_object(request, Job, request.GET.get('job_id'))
        output = {}

        appointments = Appointment.objects.filter(job=job, date__lte=datetime.date.today()) \
            .values('client') \
            .annotate(appCount=Count('id'))

        for a in appointments:
            output[a['client']] = a['appCount']

        return JsonResponse(output)

@method_decorator(csrf_exempt, name='dispatch')
class DropAppointmentsRange(CustomAPIView):
    permission_classes = [IsAuthenticated & IsOwner & IsProviderPro]

    def post(self, request):
        job = self.get_object(request, Job, request.POST.get('job_id'))
        since = request.POST.get('since')
        to = request.POST.get('to')

        # TODO notify
        appointments = Appointment.objects.filter(job=job, date__range=[since, to])
        for a in appointments:
            a.delete()

        return self.returnOK()

@method_decorator(csrf_exempt, name='dispatch')
class ScheduleTemplateView(CustomAPIView):

    def get(self, request):
        output = []
        schedule_templates = ScheduleTemplate.objects.all()
        for x in schedule_templates:
            output.append(ScheduleTemplateSerializer(x).data)
        return JsonResponse({'schedule_templates': output})


import json
import os
import random
import shutil
import time
from abc import abstractmethod
from pathlib import Path

import accelerate
import json5
import numpy as np
import torch
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration
from torch.utils.data import ConcatDataset, DataLoader
from tqdm import tqdm

from models.base.base_sampler import build_samplers
from optimizer.optimizers import NoamLR

class BaseTrainer(object):
    r"""The base trainer for all tasks. Any trainer should inherit from this class."""

    def __init__(self, args=None, cfg=None):
        super().__init__()

        self.args = args
        self.cfg = cfg

        cfg.exp_name = args.exp_name

        # init with accelerate
        self._init_accelerator()
        self.accelerator.wait_for_everyone()

        # Use accelerate logger for distributed training
        with self.accelerator.main_process_first():
            self.logger = get_logger(args.exp_name, log_level=args.log_level)

        # Log some info
        self.logger.info("=" * 56)
        self.logger.info("||\t\t" + "New training process started." + "\t\t||")
        self.logger.info("=" * 56)
        self.logger.info("\n")
        self.logger.debug(f"Using {args.log_level.upper()} logging level.")
        self.logger.info(f"Experiment name: {args.exp_name}")
        self.logger.info(f"Experiment directory: {self.exp_dir}")
        self.checkpoint_dir = os.path.join(self.exp_dir, "checkpoint")
        if self.accelerator.is_main_process:
            os.makedirs(self.checkpoint_dir, exist_ok=True)
        self.logger.debug(f"Checkpoint directory: {self.checkpoint_dir}")

        # init counts
        self.batch_count: int = 0
        self.step: int = 0
        self.epoch: int = 0
        self.max_epoch = (
            self.cfg.train.max_epoch if self.cfg.train.max_epoch > 0 else float("inf")
        )
        self.logger.info(
            "Max epoch: {}".format(
                self.max_epoch if self.max_epoch < float("inf") else "Unlimited"
            )
        )

        # Check values
        if self.accelerator.is_main_process:
            self.__check_basic_configs()
            # Set runtime configs
            self.save_checkpoint_stride = self.cfg.train.save_checkpoint_stride
            self.checkpoints_path = [
                [] for _ in range(len(self.save_checkpoint_stride))
            ]
            self.keep_last = [
                i if i > 0 else float("inf") for i in self.cfg.train.keep_last
            ]
            self.run_eval = self.cfg.train.run_eval

        # set random seed
        with self.accelerator.main_process_first():
            start = time.monotonic_ns()
            self._set_random_seed(self.cfg.train.random_seed)
            end = time.monotonic_ns()
            self.logger.debug(
                f"Setting random seed done in {(end - start) / 1e6:.2f}ms"
            )
            self.logger.debug(f"Random seed: {self.cfg.train.random_seed}")

        # setup data_loader
        with self.accelerator.main_process_first():
            self.logger.info("Building dataset...")
            start = time.monotonic_ns()
            self.train_dataloader, self.valid_dataloader = self._build_dataloader()
            end = time.monotonic_ns()
            self.logger.info(f"Building dataset done in {(end - start) / 1e6:.2f}ms")

        # setup model
        with self.accelerator.main_process_first():
            self.logger.info("Building model...")
            start = time.monotonic_ns()
            self.model = self._build_model()
            end = time.monotonic_ns()
            self.logger.debug(self.model)
            self.logger.info(f"Building model done in {(end - start) / 1e6:.2f}ms")
            self.logger.info(
                f"Model parameters: {self.__count_parameters(self.model)/1e6:.2f}M"
            )
        # optimizer & scheduler
        with self.accelerator.main_process_first():
            self.logger.info("Building optimizer and scheduler...")
            start = time.monotonic_ns()
            self.optimizer = self.__build_optimizer()
            self.scheduler = self.__build_scheduler()
            end = time.monotonic_ns()
            self.logger.info(
                f"Building optimizer and scheduler done in {(end - start) / 1e6:.2f}ms"
            )

        # accelerate prepare
        self.logger.info("Initializing accelerate...")
        start = time.monotonic_ns()
        (
            self.train_dataloader,
            self.valid_dataloader,
            self.model,
            self.optimizer,
            self.scheduler,
        ) = self.accelerator.prepare(
            self.train_dataloader,
            self.valid_dataloader,
            self.model,
            self.optimizer,
            self.scheduler,
        )
        end = time.monotonic_ns()
        self.logger.info(f"Initializing accelerate done in {(end - start) / 1e6:.2f}ms")

        # create criterion
        with self.accelerator.main_process_first():
            self.logger.info("Building criterion...")
            start = time.monotonic_ns()
            self.criterion = self._build_criterion()
            end = time.monotonic_ns()
            self.logger.info(f"Building criterion done in {(end - start) / 1e6:.2f}ms")

        # Resume or Finetune
        with self.accelerator.main_process_first():
            if args.resume:
                ## Automatically resume according to the current exprimental name
                self.logger.info("Resuming from {}...".format(self.checkpoint_dir))
                start = time.monotonic_ns()
                ckpt_path = self.__load_model(
                    checkpoint_dir=self.checkpoint_dir, resume_type=args.resume_type
                )
                end = time.monotonic_ns()
                self.logger.info(
                    f"Resuming from checkpoint done in {(end - start) / 1e6:.2f}ms"
                )
                self.checkpoints_path = json.load(
                    open(os.path.join(ckpt_path, "ckpts.json"), "r")
                )
            elif args.resume_from_ckpt_path and args.resume_from_ckpt_path != "":
                ## Resume from the given checkpoint path
                if not os.path.exists(args.resume_from_ckpt_path):
                    raise ValueError(
                        "[Error] The resumed checkpoint path {} don't exist.".format(
                            args.resume_from_ckpt_path
                        )
                    )

                self.logger.info(
                    "Resuming from {}...".format(args.resume_from_ckpt_path)
                )
                start = time.monotonic_ns()
                ckpt_path = self.__load_model(
                    checkpoint_path=args.resume_from_ckpt_path,
                    resume_type=args.resume_type,
                )
                end = time.monotonic_ns()
                self.logger.info(
                    f"Resuming from checkpoint done in {(end - start) / 1e6:.2f}ms"
                )

        # save config file path
        self.config_save_path = os.path.join(self.exp_dir, "args.json")

    ### Following are abstract methods that should be implemented in child classes ###
    @abstractmethod
    def _build_dataset(self):
        r"""Build dataset for model training/validating/evaluating."""
        pass

    @staticmethod
    @abstractmethod
    def _build_criterion():
        r"""Build criterion function for model loss calculation."""
        pass

    @abstractmethod
    def _build_model(self):
        r"""Build model for training/validating/evaluating."""
        pass

    @abstractmethod
    def _forward_step(self, batch):
        r"""One forward step of the neural network. This abstract method is trying to
        unify ``_train_step`` and ``_valid_step`` and avoid redundant implementation.
        However, for special case that using different forward step pattern for
        training and validating, you could just override this method with ``pass`` and
        implement ``_train_step`` and ``_valid_step`` separately.
        """
        pass

    @abstractmethod
    def _save_auxiliary_states(self):
        r"""To save some auxiliary states when saving model's ckpt"""
        pass

    ### Abstract methods end ###

    ### THIS IS MAIN ENTRY ###
    def train_loop(self):
        r"""Training loop. The public entry of training process."""
        # Wait everyone to prepare before we move on
        self.accelerator.wait_for_everyone()
        # dump config file
        if self.accelerator.is_main_process:
            self.__dump_cfg(self.config_save_path)
        self.model.train()
        self.optimizer.zero_grad()
        # Wait to ensure good to go
        self.accelerator.wait_for_everyone()
        while self.epoch < self.max_epoch:
            self.logger.info("\n")
            self.logger.info("-" * 32)
            self.logger.info("Epoch {}: ".format(self.epoch))

            ### TODO: change the return values of _train_epoch() to a loss dict, or (total_loss, loss_dict)
            ### It's inconvenient for the model with multiple losses
            # Do training & validating epoch
            train_loss = self._train_epoch()
            self.logger.info("  |- Train/Loss: {:.6f}".format(train_loss))
            valid_loss = self._valid_epoch()
            self.logger.info("  |- Valid/Loss: {:.6f}".format(valid_loss))
            self.accelerator.log(
                {"Epoch/Train Loss": train_loss, "Epoch/Valid Loss": valid_loss},
                step=self.epoch,
            )

            self.accelerator.wait_for_everyone()
            # TODO: what is scheduler?
            self.scheduler.step(valid_loss)  # FIXME: use epoch track correct?

            # Check if hit save_checkpoint_stride and run_eval
            run_eval = False
            if self.accelerator.is_main_process:
                save_checkpoint = False
                hit_dix = []
                for i, num in enumerate(self.save_checkpoint_stride):
                    if self.epoch % num == 0:
                        save_checkpoint = True
                        hit_dix.append(i)
                        run_eval |= self.run_eval[i]

            self.accelerator.wait_for_everyone()
            if self.accelerator.is_main_process and save_checkpoint:
                path = os.path.join(
                    self.checkpoint_dir,
                    "epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
                        self.epoch, self.step, train_loss
                    ),
                )
                self.tmp_checkpoint_save_path = path
                self.accelerator.save_state(path)
                print(f"save checkpoint in {path}")
                json.dump(
                    self.checkpoints_path,
                    open(os.path.join(path, "ckpts.json"), "w"),
                    ensure_ascii=False,
                    indent=4,
                )
                self._save_auxiliary_states()

                # Remove old checkpoints
                to_remove = []
                for idx in hit_dix:
                    self.checkpoints_path[idx].append(path)
                    while len(self.checkpoints_path[idx]) > self.keep_last[idx]:
                        to_remove.append((idx, self.checkpoints_path[idx].pop(0)))

                # Search conflicts
                total = set()
                for i in self.checkpoints_path:
                    total |= set(i)
                do_remove = set()
                for idx, path in to_remove[::-1]:
                    if path in total:
                        self.checkpoints_path[idx].insert(0, path)
                    else:
                        do_remove.add(path)

                # Remove old checkpoints
                for path in do_remove:
                    shutil.rmtree(path, ignore_errors=True)
                    self.logger.debug(f"Remove old checkpoint: {path}")

            self.accelerator.wait_for_everyone()
            if run_eval:
                # TODO: run evaluation
                pass

            # Update info for each epoch
            self.epoch += 1

        # Finish training and save final checkpoint
        self.accelerator.wait_for_everyone()
        if self.accelerator.is_main_process:
            self.accelerator.save_state(
                os.path.join(
                    self.checkpoint_dir,
                    "final_epoch-{:04d}_step-{:07d}_loss-{:.6f}".format(
                        self.epoch, self.step, valid_loss
                    ),
                )
            )
            self._save_auxiliary_states()

        self.accelerator.end_training()

    ### Following are methods that can be used directly in child classes ###
    def _train_epoch(self):
        r"""Training epoch. Should return average loss of a batch (sample) over
        one epoch. See ``train_loop`` for usage.
        """
        self.model.train()
        epoch_sum_loss: float = 0.0
        epoch_step: int = 0
        for batch in tqdm(
            self.train_dataloader,
            desc=f"Training Epoch {self.epoch}",
            unit="batch",
            colour="GREEN",
            leave=False,
            dynamic_ncols=True,
            smoothing=0.04,
            disable=not self.accelerator.is_main_process,
        ):
            # Do training step and BP
            with self.accelerator.accumulate(self.model):
                loss = self._train_step(batch)
                self.accelerator.backward(loss)
                self.optimizer.step()
                self.optimizer.zero_grad()
            self.batch_count += 1

            # Update info for each step
            # TODO: step means BP counts or batch counts?
            if self.batch_count % self.cfg.train.gradient_accumulation_step == 0:
                epoch_sum_loss += loss
                self.accelerator.log(
                    {
                        "Step/Train Loss": loss,
                        "Step/Learning Rate": self.optimizer.param_groups[0]["lr"],
                    },
                    step=self.step,
                )
                self.step += 1
                epoch_step += 1

        self.accelerator.wait_for_everyone()
        return (
            epoch_sum_loss
            / len(self.train_dataloader)
            * self.cfg.train.gradient_accumulation_step
        )

    @torch.inference_mode()
    def _valid_epoch(self):
        r"""Testing epoch. Should return average loss of a batch (sample) over
        one epoch. See ``train_loop`` for usage.
        """
        self.model.eval()
        epoch_sum_loss = 0.0
        for batch in tqdm(
            self.valid_dataloader,
            desc=f"Validating Epoch {self.epoch}",
            unit="batch",
            colour="GREEN",
            leave=False,
            dynamic_ncols=True,
            smoothing=0.04,
            disable=not self.accelerator.is_main_process,
        ):
            batch_loss = self._valid_step(batch)
            epoch_sum_loss += batch_loss.item()

        self.accelerator.wait_for_everyone()
        return epoch_sum_loss / len(self.valid_dataloader)

    def _train_step(self, batch):
        r"""Training forward step. Should return average loss of a sample over
        one batch. Provoke ``_forward_step`` is recommended except for special case.
        See ``_train_epoch`` for usage.
        """
        return self._forward_step(batch)

    @torch.inference_mode()
    def _valid_step(self, batch):
        r"""Testing forward step. Should return average loss of a sample over
        one batch. Provoke ``_forward_step`` is recommended except for special case.
        See ``_test_epoch`` for usage.
        """
        return self._forward_step(batch)

    def __load_model(
        self,
        checkpoint_dir: str = None,
        checkpoint_path: str = None,
        resume_type: str = "",
    ):
        r"""Load model from checkpoint. If checkpoint_path is None, it will
        load the latest checkpoint in checkpoint_dir. If checkpoint_path is not
        None, it will load the checkpoint specified by checkpoint_path. **Only use this
        method after** ``accelerator.prepare()``.
        """
        if checkpoint_path is None:
            ls = [str(i) for i in Path(checkpoint_dir).glob("*")]
            ls.sort(key=lambda x: int(x.split("_")[-3].split("-")[-1]), reverse=True)
            checkpoint_path = ls[0]
            self.logger.info("Resume from {}...".format(checkpoint_path))

        if resume_type in ["resume", ""]:
            # Load all the things, including model weights, optimizer, scheduler, and random states.
            self.accelerator.load_state(input_dir=checkpoint_path)

            # set epoch and step
            self.epoch = int(checkpoint_path.split("_")[-3].split("-")[-1]) + 1
            self.step = int(checkpoint_path.split("_")[-2].split("-")[-1]) + 1

        elif resume_type == "finetune":
            # Load only the model weights
            accelerate.load_checkpoint_and_dispatch(
                self.accelerator.unwrap_model(self.model),
                os.path.join(checkpoint_path, "pytorch_model.bin"),
            )
            self.logger.info("Load model weights for finetune...")

        else:
            raise ValueError("Resume_type must be `resume` or `finetune`.")

        return checkpoint_path

    # TODO: LEGACY CODE
    def _build_dataloader(self):
        Dataset, Collator = self._build_dataset()

        # build dataset instance for each dataset and combine them by ConcatDataset
        datasets_list = []
        for dataset in self.cfg.dataset:
            subdataset = Dataset(self.cfg, dataset, is_valid=False)
            datasets_list.append(subdataset)
        train_dataset = ConcatDataset(datasets_list)
        train_collate = Collator(self.cfg)
        _, batch_sampler = build_samplers(train_dataset, self.cfg, self.logger, "train")
        self.logger.debug(f"train batch_sampler: {list(batch_sampler)}")
        self.logger.debug(f"length: {train_dataset.cumulative_sizes}")
        # TODO: use config instead of (sampler, shuffle, drop_last, batch_size)
        train_loader = DataLoader(
            train_dataset,
            collate_fn=train_collate,
            batch_sampler=batch_sampler,
            num_workers=self.cfg.train.dataloader.num_worker,
            pin_memory=self.cfg.train.dataloader.pin_memory,
        )

        # Build valid dataloader
        datasets_list = []
        for dataset in self.cfg.dataset:
            subdataset = Dataset(self.cfg, dataset, is_valid=True)
            datasets_list.append(subdataset)
        valid_dataset = ConcatDataset(datasets_list)
        valid_collate = Collator(self.cfg)
        _, batch_sampler = build_samplers(valid_dataset, self.cfg, self.logger, "valid")
        self.logger.debug(f"valid batch_sampler: {list(batch_sampler)}")
        self.logger.debug(f"length: {valid_dataset.cumulative_sizes}")
        valid_loader = DataLoader(
            valid_dataset,
            collate_fn=valid_collate,
            batch_sampler=batch_sampler,
            num_workers=self.cfg.train.dataloader.num_worker,
            pin_memory=self.cfg.train.dataloader.pin_memory,
        )
        return train_loader, valid_loader

    @staticmethod
    def _set_random_seed(seed):
        r"""Set random seed for all possible random modules."""
        random.seed(seed)
        np.random.seed(seed)
        torch.random.manual_seed(seed)

    def _check_nan(self, loss, y_pred, y_gt):
        if torch.any(torch.isnan(loss)):
            self.logger.fatal("Fatal Error: Training is down since loss has Nan!")
            self.logger.error("loss = {:.6f}".format(loss.item()), in_order=True)
            if torch.any(torch.isnan(y_pred)):
                self.logger.error(
                    f"y_pred has Nan: {torch.any(torch.isnan(y_pred))}", in_order=True
                )
            else:
                self.logger.debug(
                    f"y_pred has Nan: {torch.any(torch.isnan(y_pred))}", in_order=True
                )
            if torch.any(torch.isnan(y_gt)):
                self.logger.error(
                    f"y_gt has Nan: {torch.any(torch.isnan(y_gt))}", in_order=True
                )
            else:
                self.logger.debug(
                    f"y_gt has nan: {torch.any(torch.isnan(y_gt))}", in_order=True
                )
            if torch.any(torch.isnan(y_pred)):
                self.logger.error(f"y_pred: {y_pred}", in_order=True)
            else:
                self.logger.debug(f"y_pred: {y_pred}", in_order=True)
            if torch.any(torch.isnan(y_gt)):
                self.logger.error(f"y_gt: {y_gt}", in_order=True)
            else:
                self.logger.debug(f"y_gt: {y_gt}", in_order=True)

            # TODO: still OK to save tracking?
            self.accelerator.end_training()
            raise RuntimeError("Loss has Nan! See log for more info.")

    ### Protected methods end ###

    ## Following are private methods ##
    ## !!! These are inconvenient for GAN-based model training. It'd be better to move these to svc_trainer.py if needed.
    def __build_optimizer(self):
        r"""Build optimizer for model."""
        # Make case-insensitive matching
        if self.cfg.train.optimizer.lower() == "adadelta":
            optimizer = torch.optim.Adadelta(
                self.model.parameters(), **self.cfg.train.adadelta
            )
            self.logger.info("Using Adadelta optimizer.")
        elif self.cfg.train.optimizer.lower() == "adagrad":
            optimizer = torch.optim.Adagrad(
                self.model.parameters(), **self.cfg.train.adagrad
            )
            self.logger.info("Using Adagrad optimizer.")
        elif self.cfg.train.optimizer.lower() == "adam":
            optimizer = torch.optim.Adam(self.model.parameters(), **self.cfg.train.adam)
            self.logger.info("Using Adam optimizer.")
        elif self.cfg.train.optimizer.lower() == "adamw":
            optimizer = torch.optim.AdamW(
                self.model.parameters(), **self.cfg.train.adamw
            )
        elif self.cfg.train.optimizer.lower() == "sparseadam":
            optimizer = torch.optim.SparseAdam(
                self.model.parameters(), **self.cfg.train.sparseadam
            )
        elif self.cfg.train.optimizer.lower() == "adamax":
            optimizer = torch.optim.Adamax(
                self.model.parameters(), **self.cfg.train.adamax
            )
        elif self.cfg.train.optimizer.lower() == "asgd":
            optimizer = torch.optim.ASGD(self.model.parameters(), **self.cfg.train.asgd)
        elif self.cfg.train.optimizer.lower() == "lbfgs":
            optimizer = torch.optim.LBFGS(
                self.model.parameters(), **self.cfg.train.lbfgs
            )
        elif self.cfg.train.optimizer.lower() == "nadam":
            optimizer = torch.optim.NAdam(
                self.model.parameters(), **self.cfg.train.nadam
            )
        elif self.cfg.train.optimizer.lower() == "radam":
            optimizer = torch.optim.RAdam(
                self.model.parameters(), **self.cfg.train.radam
            )
        elif self.cfg.train.optimizer.lower() == "rmsprop":
            optimizer = torch.optim.RMSprop(
                self.model.parameters(), **self.cfg.train.rmsprop
            )
        elif self.cfg.train.optimizer.lower() == "rprop":
            optimizer = torch.optim.Rprop(
                self.model.parameters(), **self.cfg.train.rprop
            )
        elif self.cfg.train.optimizer.lower() == "sgd":
            optimizer = torch.optim.SGD(self.model.parameters(), **self.cfg.train.sgd)
        else:
            raise NotImplementedError(
                f"Optimizer {self.cfg.train.optimizer} not supported yet!"
            )
        return optimizer

    def __build_scheduler(self):
        r"""Build scheduler for optimizer."""
        # Make case-insensitive matching
        if self.cfg.train.scheduler.lower() == "lambdalr":
            scheduler = torch.optim.lr_scheduler.LambdaLR(
                self.optimizer, **self.cfg.train.lambdalr
            )
        elif self.cfg.train.scheduler.lower() == "multiplicativelr":
            scheduler = torch.optim.lr_scheduler.MultiplicativeLR(
                self.optimizer, **self.cfg.train.multiplicativelr
            )
        elif self.cfg.train.scheduler.lower() == "steplr":
            scheduler = torch.optim.lr_scheduler.StepLR(
                self.optimizer, **self.cfg.train.steplr
            )
        elif self.cfg.train.scheduler.lower() == "multisteplr":
            scheduler = torch.optim.lr_scheduler.MultiStepLR(
                self.optimizer, **self.cfg.train.multisteplr
            )
        elif self.cfg.train.scheduler.lower() == "constantlr":
            scheduler = torch.optim.lr_scheduler.ConstantLR(
                self.optimizer, **self.cfg.train.constantlr
            )
        elif self.cfg.train.scheduler.lower() == "linearlr":
            scheduler = torch.optim.lr_scheduler.LinearLR(
                self.optimizer, **self.cfg.train.linearlr
            )
        elif self.cfg.train.scheduler.lower() == "exponentiallr":
            scheduler = torch.optim.lr_scheduler.ExponentialLR(
                self.optimizer, **self.cfg.train.exponentiallr
            )
        elif self.cfg.train.scheduler.lower() == "polynomiallr":
            scheduler = torch.optim.lr_scheduler.PolynomialLR(
                self.optimizer, **self.cfg.train.polynomiallr
            )
        elif self.cfg.train.scheduler.lower() == "cosineannealinglr":
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
                self.optimizer, **self.cfg.train.cosineannealinglr
            )
        elif self.cfg.train.scheduler.lower() == "sequentiallr":
            scheduler = torch.optim.lr_scheduler.SequentialLR(
                self.optimizer, **self.cfg.train.sequentiallr
            )
        elif self.cfg.train.scheduler.lower() == "reducelronplateau":
            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                self.optimizer, **self.cfg.train.reducelronplateau
            )
        elif self.cfg.train.scheduler.lower() == "cycliclr":
            scheduler = torch.optim.lr_scheduler.CyclicLR(
                self.optimizer, **self.cfg.train.cycliclr
            )
        elif self.cfg.train.scheduler.lower() == "onecyclelr":
            scheduler = torch.optim.lr_scheduler.OneCycleLR(
                self.optimizer, **self.cfg.train.onecyclelr
            )
        elif self.cfg.train.scheduler.lower() == "cosineannearingwarmrestarts":
            scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
                self.optimizer, **self.cfg.train.cosineannearingwarmrestarts
            )
        elif self.cfg.train.scheduler.lower() == "noamlr":
            scheduler = NoamLR(self.optimizer, **self.cfg.train.lr_scheduler)
        else:
            raise NotImplementedError(
                f"Scheduler {self.cfg.train.scheduler} not supported yet!"
            )
        return scheduler

    def _init_accelerator(self):
        self.exp_dir = os.path.join(
            os.path.abspath(self.cfg.log_dir), self.args.exp_name
        )
        project_config = ProjectConfiguration(
            project_dir=self.exp_dir,
            logging_dir=os.path.join(self.exp_dir, "log"),
        )
        self.accelerator = accelerate.Accelerator(
            gradient_accumulation_steps=self.cfg.train.gradient_accumulation_step,
            log_with=self.cfg.train.tracker,
            project_config=project_config,
        )
        if self.accelerator.is_main_process:
            os.makedirs(project_config.project_dir, exist_ok=True)
            os.makedirs(project_config.logging_dir, exist_ok=True)
        with self.accelerator.main_process_first():
            self.accelerator.init_trackers(self.args.exp_name)

    def __check_basic_configs(self):
        if self.cfg.train.gradient_accumulation_step <= 0:
            self.logger.fatal("Invalid gradient_accumulation_step value!")
            self.logger.error(
                f"Invalid gradient_accumulation_step value: {self.cfg.train.gradient_accumulation_step}. It should be positive."
            )
            self.accelerator.end_training()
            raise ValueError(
                f"Invalid gradient_accumulation_step value: {self.cfg.train.gradient_accumulation_step}. It should be positive."
            )
        # TODO: check other values

    @staticmethod
    def __count_parameters(model):
        model_param = 0.0
        if isinstance(model, dict):
            for key, value in model.items():
                model_param += sum(p.numel() for p in model[key].parameters())
        else:
            model_param = sum(p.numel() for p in model.parameters())
        return model_param

    def __dump_cfg(self, path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        json5.dump(
            self.cfg,
            open(path, "w"),
            indent=4,
            sort_keys=True,
            ensure_ascii=False,
            quote_keys=True,
        )

    ### Private methods end ###


import os
import pefile

__LIST_OF_DLLS = {
    'ADVAP132.DLL', # Advanced Win32 application programming interfaces
    'AWFAXP32.DLL', # Mail API fax transport
    'AWFXAB32.DLL', # Address book
    'AWPWD32.DLL', # Security support
    'AWRESX32.DLL', # Resource Executor
    'AWUTIL32.DLL', # At Work Security Support
    'BHNETB.DLL', # Network monitor SMS client
    'BHSUPP.DLL', # Network monitor SMS client
    'CCAPI.DLL', # Microsoft Network component
    'CCEI.DLL', # Microsoft Network component
    'CCPSH.DLL', # Microsoft Network component
    'CCTN20.DLL', # Microsoft Network component
    'CMC.DLL', # Common messaging calls for Mail API 1.0
    'COMCTL32.DLL', # User Experience Controls Library
    'COMDLG32.DLL', # Common Dialogue Library
    'CRTDLL.DLL', # Microsoft C Runtime Library
    'DCIMAN.DLL', # Display Control Interface Manager
    'DCIMAN32.DLL', # Display Control Interface Manager
    'DSKMAINT.DLL', # Disk Utilities engine
    'GDI32.DLL', # GDI Client DLL
    'GROUP.DLL', # policy support
    'HYPERTERM.DLL', # Terminal DLL
    'KERNL32.DLL', # Windows NT BASE API Client DLL
    'LZ32.DLL', # LZ Expand/Compress API DLL
    'MAPI.DLL', # Mail / Exchange component
    'MAPI32.DLL', # Extended MAPI 1.0 for Windows NT
    'MFC30.DLL', # Shared MFC DLL
    'MPR.DLL', # Multiple Provider Router DLL
    'MSPST32.DLL', # Microsoft Personal Folder/Address Book Service Provider
    'MSFS32.DLL', # MAPI 1.0 Service Providers for Microsoft Mail
    'MSNDUI.DLL', # Microsoft Network component
    'MSNET32.DLL', # Microsoft 32-bit Network API Library
    'MSSHRUI.DLL', # Shell extensions for sharing
    'MSVIEWUT.DLL', # Service data-link libraries for display engines
    'NAL.DLL', # Network monitor SMS client
    'NDIS30.DLL', # Network monitor SMS client
    'NETAPI.DLL', # Network API
    'NETAPI32.DLL', # Net Win32 API DLL
    'NETBIOS.DLL', # NetBIOS API Library
    'NETDI.DLL', # Net Device installer
    'NETSETUP.DLL', # Network server-based setup
    'NWAB32.DLL', # Address book provider
    'NWNET32.DLL', # NetWare client
    'NWNP32.DLL', # NetWare component
    'OLEDLG.DLL', # Microsoft Windows OLE 2.0 User Interface Support
    'POWERCFG.DLL', # Advanced Power Management Control Panel
    'RASPI.DLL', # Automated Software Profile, Analysis, Removal and Signature Information
    'RASAPI16.DLL', # Remote Access Services 16-bit API Library
    'RASAPI32.DLL', # Remote Access 16-bit API Library
    'RPCRT4.DLL', # Remote Procedure Call Runtime
    'RPCLTC1.DLL', # Remote Procedure Call libraries
    'RPCTLC3.DLL', # Remote Procedure Call libraries
    'RPCTLC5.DLL', # Remote Procedure Call libraries
    'RPCTLC6.DLL', # Remote Procedure Call libraries
    'RPCTLS3.DLL', # Remote Procedure Call libraries
    'RPCTLS5.DLL', # Remote Procedure Call libraries
    'RPCTLS6.DLL', # Remote Procedure Call libraries
    'RPCNS4.DLL', # Remote Procedure Call Name Service Client
    'RSRC32.DLL', # Resource Meter
    'SAPNSP.DLL', # Winsock data-link library
    'SECUR32.DLL', # Security Support Provider Interface
    'SHELL32.DLL', # Windows Shell Common DLL
    'SLENH.DLL', # Advanced Power Management options
    'SHLWAPI.DLL', # Library for UNC and URL Paths, Registry Entries and Color Settings
    'UMDM32.DLL', # Universal Modem Driver component
    'USER32.DLL', # USER API Client DLL
    'VERSION.DLL', # Version Checking and File Installation Libraries
    'WININET.DLL', # Internet Extensions for Win32
    'WINMM.DLL', # MCI API DLL
    'WINREG.DLL', # Remote Registry support
    'WINSOCK.DLL', # Socket API for Windows
    'WS2.DLL', # 32.DLL Windows Socket 2.0 32-Bit DLL
    'WSOCK32.DLL', # Windows Socket 32-Bit DLL
}

__SET_OF_DLLS = set(__LIST_OF_DLLS)

# https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-_image_file_header
__FILE_HEADER = [
    'Machine', # The architecture type of the computer.
    'NumberOfSections', # The number of sections.
    'TimeDateStamp', # The low 32 bits of the time stamp of the image.
    'PointerToSymbolTable', # The offset of the symbol table, in bytes, or zero if no COFF symbol table exists.
    'NumberOfSymbols', # The number of symbols in the symbol table.
    'SizeOfOptionalHeader', # The size of the optional header, in bytes.
    'Characteristics' # The characteristics of the image.
]

# https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-image_optional_header32
__OPTIONAL_HEADER = [
    'Magic', # The state of the image file.
    'MajorLinkerVersion', # The major version number of the linker.
    'MinorLinkerVersion', # The minor version number of the linker.
    'SizeOfCode', # The size of the code section, in bytes, or the sum of all such sections if there are multiple code sections.
    'SizeOfInitializedData', # The size of the initialized data section, in bytes, or the sum of all such sections if there are multiple initialized data sections.
    'SizeOfUninitializedData', # The size of the uninitialized data section, in bytes, or the sum of all such sections if there are multiple uninitialized data sections.
    'AddressOfEntryPoint', # A pointer to the entry point function, relative to the image base address.
    'BaseOfCode', # A pointer to the beginning of the code section, relative to the image base.
    'BaseOfData', # A pointer to the beginning of the data section, relative to the image base.
    'ImageBase', # The preferred address of the first byte of the image when it is loaded in memory.
    'SectionAlignment', # The alignment of sections loaded in memory, in bytes.
    'FileAlignment', # The alignment of the raw data of sections in the image file, in bytes.
    'MajorOperatingSystemVersion', # The major version number of the required operating system.
    'MinorOperatingSystemVersion', # The minor version number of the required operating system.
    'MajorImageVersion', # The major version number of the image.
    'MinorImageVersion', # The minor version number of the image.
    'MajorSubsystemVersion', # The major version number of the subsystem.
    'MinorSubsystemVersion', # The minor version number of the subsystem.
    'Reserved1',  # (Win32VersionValue) This member is reserved and must be 0.
    'SizeOfImage', # The size of the image, in bytes, including all headers.
    'SizeOfHeaders', # The combined size of the following items, rounded to a multiple of the value specified in the FileAlignment member.
    'CheckSum', # The image file checksum.
    'Subsystem', # The subsystem required to run this image.
    'DllCharacteristics', # The DLL characteristics of the image.
    'SizeOfStackReserve', # The number of bytes to reserve for the stack.
    'SizeOfStackCommit', # The number of bytes to commit for the stack.
    'SizeOfHeapReserve', # The number of bytes to commit for the local heap.
    'SizeOfHeapCommit', # This member is obsolete.
    'LoaderFlags', # The number of directory entries in the remainder of the optional header.
    'NumberOfRvaAndSizes' # A pointer to the first IMAGE_DATA_DIRECTORY structure in the data directory.
]

__DIRECTORY_ENTRY_TYPES = {
    'IMAGE_DIRECTORY_ENTRY_EXPORT' : 0,
    'IMAGE_DIRECTORY_ENTRY_IMPORT' : 1,
    'IMAGE_DIRECTORY_ENTRY_RESOURCE' : 2,
    'IMAGE_DIRECTORY_ENTRY_EXCEPTION' : 3,
    'IMAGE_DIRECTORY_ENTRY_SECURITY' : 4,
    'IMAGE_DIRECTORY_ENTRY_BASERELOC' : 5,
    'IMAGE_DIRECTORY_ENTRY_DEBUG' : 6,
    'IMAGE_DIRECTORY_ENTRY_COPYRIGHT' : 7,
    'IMAGE_DIRECTORY_ENTRY_GLOBALPTR' : 8,
    'IMAGE_DIRECTORY_ENTRY_TLS' : 9,
    'IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG' : 10,
    'IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT' : 11,
    'IMAGE_DIRECTORY_ENTRY_IAT' : 12,
    'IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT' : 13,
    'IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR' : 14,
    'IMAGE_DIRECTORY_ENTRY_RESERVED' : 15
}
# https://docs.microsoft.com/en-us/windows/win32/debug/pe-format#section-table-section-headers
__SECTION_HEADER = [
    'VirtualSize',
    'VirtualAddress',
    'SizeOfRawData',
    'PointerToRawData',
    'PointerToRelocations',
    'PointerToLinenumbers',
    'NumberOfRelocations',
    'NumberOfLinenumbers',
    'Characteristics'
]

__RESOURCE_DIRECTORY_TABLE = [
    'Characteristics',
    'MajorVersion',
    'MinorVersion',
    'NumberOfIdEntries',
    'NumberOfNamedEntries',
    'TimeDateStamp'
]

def extract_feature(file_path):
    with open(file_path, 'rb') as f:
        file_data = f.read()
    try:
        pe = pefile.PE(data=file_data)
        ret = dict()
        # Feature From DLLs referred
        for dll in __LIST_OF_DLLS:
            ret[dll] = 0
        if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
            for lib in pe.DIRECTORY_ENTRY_IMPORT:
                if lib.dll != None:
                    if lib.dll.decode().upper() in __SET_OF_DLLS :
                        ret[lib.dll.decode().upper()] = 1
        # Feature From COFF file header
        for member in __FILE_HEADER:
            ret[member] = getattr(pe.FILE_HEADER, member, -1)
        # Feature From Optional header: standard fields and Windows specific fields
        for member in __OPTIONAL_HEADER:
            ret[member] = getattr(pe.OPTIONAL_HEADER, member, -1)
        # Feature From Optional header: data directories
        for key in __DIRECTORY_ENTRY_TYPES.keys():
            ret[f'{key}:VirtualAddress'] = -1
            ret[f'{key}:Size'] = -1
        for structure in pe.OPTIONAL_HEADER.DATA_DIRECTORY:
            ret[f'{structure.name}:VirtualAddress'] = structure.VirtualAddress
            ret[f'{structure.name}:Size'] = structure.Size
        # Feature From Section headers
        for section_name in ['text', 'data', 'resource']:
            for member_name in __SECTION_HEADER :
                ret[f'{section_name}:{member_name}'] = -1
        for section in pe.sections:
            scn = section.Name
            if section.Name == b'.text\x00\x00\x00':
                for member in __SECTION_HEADER:
                    ret[f'text:{member}'] = getattr(section, member, -1)
            elif section.Name == b'.data\x00\x00\x00':
                for member in __SECTION_HEADER:
                    ret[f'data:{member}'] = getattr(section, member, -1)
            elif section.Name == b'.rsrc\x00\x00\x00':
                for member in __SECTION_HEADER:
                    ret[f'resource:{member}'] = getattr(section, member, -1)

        # Feature From Resource directory table & resources
        for resource_type in ['Cursors', 'Bitmaps', 'Icons', 'Menus', 'Dialogs', 'Strings', 'Fonts', 'Group Cursors', 'Group Icons']:
            ret[resource_type] = 0
        for resource_struct in __RESOURCE_DIRECTORY_TABLE:
            ret[f'Resource:{resource_struct}'] = -1
        if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
            for resource_struct in __RESOURCE_DIRECTORY_TABLE:
                ret[f'Resource:{resource_struct}'] = getattr(pe.DIRECTORY_ENTRY_RESOURCE.struct, resource_struct, -1)
            for resource_type in pe.DIRECTORY_ENTRY_RESOURCE.entries:
                # Cursors
                if resource_type.struct.Id == 1:
                    ret['Cursors'] += len(resource_type.directory.entries)
                # Bitmaps
                elif resource_type.struct.Id == 2:
                    ret['Bitmaps'] += len(resource_type.directory.entries)
                # Icons
                elif resource_type.struct.Id == 3:
                    ret['Icons'] += len(resource_type.directory.entries)
                # Menus
                elif resource_type.struct.Id == 4:
                    ret['Menus'] += len(resource_type.directory.entries)
                # Dialogs
                elif resource_type.struct.Id == 5 :
                    ret['Dialogs'] += len(resource_type.directory.entries)
                # Strings
                elif resource_type.struct.Id == 6:
                    ret['Strings'] += len(resource_type.directory.entries)
                # Fonts
                elif resource_type.struct.Id == 8:
                    ret['Fonts'] += len(resource_type.directory.entries)
                # Group Cursors
                elif resource_type.struct.Id == 12:
                    ret['Group Cursors'] += len(resource_type.directory.entries)
                # Group Icons
                elif resource_type.struct.Id == 14:
                    ret['Group Icons'] += len(resource_type.directory.entries)
        return ret
    except pefile.PEFormatError as pefe:
        print("PEFormatError", pefe, os.path.basename(file_path))
        return None
    except Exception as e:
        print("Exception", e, os.path.basename(file_path))
        return None


#!/home/ocean/anaconda3/bin/python3
from numpy import cos, arccos, sin, arctan, tan, pi, sqrt; from numpy import array as ary; import numpy as np; tau = 2*pi
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
import seaborn as sns
#fluxFit.py only deals with reading data and then fitting them.
from handleFit import *
#Oh my god the order matters:
#if you import handleFit AFTER flusDist, then you'll get into a loopy mess.
#But the current configuration works, i.e. BEFORE.
from fluxDists import *	#import the models to fit the data against.
from HiddenSubprograms import *	#other lower level programs
from matplotlib.ticker import FormatStrFormatter
import matplotlib.ticker as ticker
#print("fluxFit",__name__)
global FastFlux,fastFErr, ThmFlux,ThmErr, τ
τ = 368	#unit:cm^2

target_y=0	#picking the xz plane that contains the source.
print("for y=",target_y)
x,z, AllFlux, allFErr = readData("n-Diff-new.csv",y_val=target_y)
x2,z2,FastFlux,fastFErr=readData("CadmiumNew.csv",y_val=target_y)
xzfast = ary([x2,z2]).T#of shape (36,2)
xzThm,ThmFlux,ThmErr = Extended_thermal_flux(x,z,AllFlux,allFErr,x2,z2,FastFlux,fastFErr)
def zThreshold(llim=0, ulim= 200):
	global xzfast,FastFlux,fastFErr, xzThm,ThmFlux,ThmErr
	fast_1dpl,fast_0dpl=101.25,141.75
	xzThm,ThmFlux,ThmErr = RemoveData(xzThm,ThmFlux,ThmErr,llim,ulim)
	if (llim<=fast_0dpl):
		xzfast,FastFlux,fastFErr=RemoveData(xzfast,FastFlux,fastFErr,llim, ulim)
		if (llim>fast_1dpl):
			#raise ImportWarning("only one datapoint parsed for fast")
			print("Warning: only one data point parsed for fast neutron flux.")
	return
def xThreshold(llim=-61, ulim=61):
	global xzfast,FastFlux,fastFErr, xzThm,ThmFlux,ThmErr
	xzThm,ThmFlux,ThmErr	=RemoveDatax(xzThm,ThmFlux,ThmErr,llim,ulim)
	xzfast,FastFlux,fastFErr=RemoveDatax(xzfast,FastFlux,fastFErr,llim, ulim)
	return
def zT(ind):
	Slice = [(20.0,21.0),(40.0,41.0),(60.0,61.0), (80.0,82.0),(101.0,102.0),(141.0,142.0),(182.0,183.0)]
	low,upp=  Slice[ind]
	zThreshold(low,upp)
def xT(ind):
	Slice = [(-61.0,-60.0),(-41.0,-40.0),(-21.0,-20.0),(-1.0,1.0),(20.0,21.0),(40.0,41.0),(60.0,61.0)]
	low,upp=  Slice[ind]
	xThreshold(low,upp)
def zCutoff(ind):
	if ind>=0:
		zThreshold((20.0,40.0,60.0,80.0)[ind])
	elif ind<0:
		zThreshold(llim=0,ulim=(25.0,45.0,65.0,85.0)[-ind])
def scaleUpError(SF):
	global fastFErr, ThmErr
	fastFErr, ThmErr = FastFlux*SF, ThmErr*SF
	return
def SquareError():
	global fastFErr, ThmErr
	fastFErr, ThmErr = FastFlux**2, ThmErr**2
	return
#_______________________________________________CHANGE THE FOLLOWING
#zT(0)
#xT(3)
zCutoff(1)
SF=.70
#SquareError()
#scaleUpError(10);SquareError()#;scaleUpError(5)
#_______________________________________________END of changeable stuff

if __name__=="__main__":
	smoothData, plotObject=handleFittingResult(
	FastSquare,  xzfast,FastFlux,fastFErr, [368,77]);				titleName="Fast neutron flux distribution"
	#FastSquare_τ,xzfast,FastFlux,fastFErr, [77]);					titleName="Fast neutron flux distribution"
	#FastSph,     xzfast,FastFlux,fastFErr, [368]);					titleName="Fast neutron flux distribution \nwith spherically symmetric fit"
	#FastSph_τ,   xzfast,FastFlux,fastFErr,[]);						titleName="Fast neutron flux distribution \nwith spherically symmetric fit"
	#OneGrp_Thm, xzThm,ThmFlux,ThmErr,	[60,77]);					titleName="Thermal neutron flux distribution \nwith 1 group fit"
	#TwoGrp_Thm, xzThm,ThmFlux,ThmErr,	[60,77]);					titleName="Thermal neutron flux distribution \nwith 2 group fit with fixed τ"
	#TwoGrp_redefine_τ,xzThm,ThmFlux,ThmErr, [60,368,77]);			titleName="Thermal neutron flux distribution \nwith 2 group fit with"
	#OneGrp_Thm_Sph,xzThm,ThmFlux,ThmErr,[160]);					titleName="Thermal neutron flux distribution \nwith 1 group spherically symmetric fit"
	#ThirdOrd_OneGrp_Thmz,xzThm,ThmFlux,ThmErr,[160,0.4,77]);		titleName="Thermal neutron flux distribution \n1 group fit with 3rd order term"
	#ThirdOrd_TwoGrp_Thmz,xzThm,ThmFlux,ThmErr,[160,368,0.4,77]);	titleName="Thermal neutron flux distribution \n2 group fit with 3rd order term"
	#FixL,		xzThm,	ThmFlux,	ThmErr,	[0.2]);					titleName="Thermal neutron flux distribution \n1 group fit with 3rd order term"
	#FixLtau,	xzThm,	ThmFlux,	ThmErr,	[0.2]);					titleName="Thermal neutron flux distribution \n2 group fit with 3rd order term"
	#OneGrp_ThmNoBD, xzThm, ThmFlux, ThmErr, [55,20]);				titleName="Thermal neutron flux distribution \nwith 1 group fit"
	#GreensFuncFast,	xzfast, FastFlux,fastFErr, [76,368,121]);	titleName="Sinusoidal Summation for\n fast neutron flux distribution"
	#GreensFunc,		xzThm, ThmFlux, ThmErr, [76,368,121,60]);	titleName="Sinusoidal Summation for\n thermal neutron flux distribution"
	titleName+=' at y='+str(target_y)+" cm"
	fileName=abbrName(titleName)
	FAST =(("Fast" in fileName) or ("fast" in titleName))
	if FAST:
		heatmapDF = ToDataFrame(xzfast,FastFlux,fastFErr)
	else:
		heatmapDF = ToDataFrame(xzThm, ThmFlux, ThmErr  )
	x,z,flux,err = heatmapDF.values.T #resuing the x,z variables for more important things

	print( "Normalization constant=", f(readA()) )
	smoothData = pd.DataFrame(smoothData,columns=["x","z","flux"])
	if (heatmapDF.nunique()["z"]>1) and (heatmapDF.nunique()["x"]>1):
		#plotObject is a graph itself
		(plotObject,equation) = plotObject
		#set the title
		ttl_txt = plotObject.fig.suptitle("Residuals for fit "+get_expression())#,fontsize=20)#,xy=[.15,.8])
		if "set figure size"=="set figure size":
			fig_size = (plotObject.fig.get_size_inches())
			fig_size = fig_size*SF
			plotObject.fig.set_size_inches(fig_size)
		if heatmapDF.min()["z"]!=20.25:
			plotObject.savefig("Printable/"+fileName+" HeatMap with cutoff of z≥"+str(heatmapDF.min()["z"])+" cm RESIDUAL.png",dpi=100)
		else:
			plotObject.savefig("Printable/"+fileName+" HeatMap RESIDUAL.png",dpi=100)
		if "to clear figure"=="to clear figure":	#indented for clarity
			plt.cla()
			ttl_txt.set_text("")
			equation.set_text("")
		ax = plt.subplot(111)
		#plot the actual fit (in heat map form)
		smoothflux = smoothData.pivot("z","x","flux")
		smooth_x, smooth_z = np.meshgrid(smoothflux.columns.values,smoothflux.index)
		#Two cmaps to choose from, both from seaborn:
		#cmap = sns.diverging_palette(250, 15, s=75, l=40, center="dark",as_cmap=True)
		if not FAST:	#thermal neutron color palette
			cmap = sns.blend_palette(["b","r"],as_cmap=True)
		elif FAST:
			cmap = sns.blend_palette(["g","y"],as_cmap=True)
		heat = ax.pcolor(smooth_x,smooth_z,smoothflux, cmap=cmap)
		scat = ax.scatter(x,z,c=flux,cmap=cmap,edgecolor="w")
		cbar = plt.colorbar(heat,drawedges=False,ax=ax)
		if "formatting"=="formatting":# can remove anytime if I don't like the formatting:
			ax.set_xlabel("x (cm)")
			ax.set_ylabel("z (cm)")
			ax.axis([x.min()-1, x.max()+1, z.min()-1, z.max()+1])
			plt.tight_layout()
			for spine in ax.spines.values():
				spine.set_visible(False)
			cbar.outline.set_visible(False)
			cbar.set_label("counts/s")
			plt.figtext(0.5,0.001,"Background =best fit; colored dots = actual data.",wrap=True,ha="center")
			ax.set_title(titleName,va="top")
		if heatmapDF.min()["z"]!=20.25:
			fileName = "Printable/"+fileName+" HeatMap with cutoff of z≥"+str(heatmapDF.min()["z"])+" cm.png"
		else:
			fileName = "Printable/"+fileName+" HeatMap.png"
		plt.rcParams["figure.figsize"] = fig_size
		#plt.set_size_inches(fig_size)
		plt.savefig(fileName,dpi=100)
	elif (heatmapDF.nunique()["z"]==1):
		smoothData = smoothData.drop("z",axis=1)
		#smoothx,smoothy = smoothData.drop("z",axis=1).values.T
		fig,(fit,resid) = plt.subplots(2,1, gridspec_kw={'height_ratios':[3,2]})
		resid = residPlot(plotObject,resid)	#plot the residual
		smoothData.plot(x="x", y="flux" ,label = get_expression() , ax=fit, sharex=resid)
		#must unpack like this because DataFrame doesn't allow errorbar plot directly.
		z_str = f(heatmapDF["z"][0])
		fit.errorbar(x,flux,err, label="z="+z_str+"cm", fmt="o",ecolor="g",capsize=4, markersize=4)
		#take care of the auxillary stuff
		fit.legend()
		fit.set_ylabel(r"count rate ($s^{-1}$)")
		fit.set_title(titleName,)#va="top")
		fit.set_ylim(0)
		if "scaling"=="scaling":
			fig_size = fig.get_size_inches()
			fig_size = fig_size*SF
			fig.set_size_inches(fig_size)
		fileName="Printable/"+fileName+" z="+z_str+".png"
		fig.savefig(fileName)
	else:
		assert heatmapDF.nunique()["x"]==1
		smoothData = smoothData.drop("x",axis=1)
		#smoothx,smoothy = smoothData.drop("z",axis=1).values.T
		fig,(fit,resid) = plt.subplots(2,1, gridspec_kw={'height_ratios':[3,2]})
		resid = residPlot(plotObject,resid)	#plot the residual
		smoothData.plot(x="z", y="flux" ,label = get_expression() , ax=fit, sharex=resid)
		x_str = f(heatmapDF["x"][0])
		fit.errorbar(z,flux,err, label="x="+x_str+"cm", fmt="o",ecolor="g",capsize=4, markersize=4)
		fit.legend()
		fit.set_ylabel(r"count rate ($s^{-1}$)")
		fit.set_title(titleName,)#va='top')
		fit.set_ylim(0)
		resid.set_xlabel("z(cm)")
		if "scaling"=="scaling":
			fig_size = fig.get_size_inches()
			fig_size = fig_size*SF
			fig.set_size_inches(fig_size)
		fileName= "Printable/"+(fileName)+" x="+x_str+".png"
		fig.savefig(fileName)
	print("fileName=",fileName)
	'''
	bugs to fix:
			2.0 make sure smooth heat map is in the correct direction, may have to tweak the [::-1] in mymesh	#DONE
			2. Put the actual data on top of heat map	#DONE
			2.1 using	#DONE
			ax.scatter()	#DONE
			2.2 Ensure that there are no text overlaps
			#2.3 Remove the text and title that has remained in the background of ax
			#Fixed using https://stackoverflow.com/questions/10559144/matplotlib-suptitle-prints-over-old-title
			2.3 format the ytick and xtick.	#DONE
			1. Make the program work for only a single slice of z data as well.
			1.0 Make sure mymesh works with zero as well.#DONE
			1.1 make sure that it prints out one graph (i.e. the following line should work:)
			1.2 with no overlap* text, and good layout	#DONE
			1.3 Legend should be readable	#DONE
	In fluxDists:
			1.6 Add units on cbar 	#DONE
			1.7 add captions using	#DONE
			1.4 make another batch of models that have exponential part replaced by 1	#DONE
			1.5 And redirect the current models to those ones when they detect that len(z.unique())==1	#DONE
			1.9 Carefully look through every single text expression for the fit equation used.
			1.9.2But where is that implemented???
			1.9.3 make an equivalent method of xthreshold; and deal/w /understand line 62 in handleFit.py #DONE
	1.8 Think about the title+their placements
	4. Make the actual plot, using a color scheme "b/r" for thermal, and a color scheme "g/y" for fast
	'''
else:
	print("End of importing! Weeee_______________________________________")


# -*- coding: utf-8 -*-
"""
Created on Mon Apr  6 20:43:33 2020

@author: iant

MAKE LNO MOLECULAR SPECTRA FOR A GIVEN ORDER/WAVENUMBER RANGE
"""
import numpy as np
import os
import matplotlib.pyplot as plt

from repos.pytran.pytran.hitran_utils import get_molecule_id
from repos.pytran.pytran.hitran import read_hitran2012_parfile, calculate_hitran_xsec
from repos.nomad_tools.nomadtools.rcsetup import defaultParams

def get_xsection(molecule, nu_hr, Smin=0.0, temperature=210.0, pressure=0.0):

    nu_hr_min = np.min(nu_hr)
    nu_hr_max = np.max(nu_hr)

    M = get_molecule_id(molecule)
    filename = os.path.join(defaultParams['paths.dirLP'], '%02i_hit16_2000-5000_CO2broadened.par' % M)
    if not os.path.exists(filename):
        filename = os.path.join(defaultParams['paths.dirLP'], '%02i_hit16_2000-5000.par' % M)

    if Smin == 0.0:
        if molecule == "CO2":
            Smin=1.0e-26
        elif molecule == "H2O":
            Smin=1.0e-27
        else:
            Smin=1.0e-33

    nlines = 999
    Smin /= 10.0
    while nlines>200:
        Smin *= 10.0
        LineList = read_hitran2012_parfile(filename, nu_hr_min, nu_hr_max, Smin=Smin)
        nlines = len(LineList['S'])
        print('Found %i lines for %s Smin of %0.1g' %(nlines, molecule, Smin))
    return calculate_hitran_xsec(LineList, M, nu_hr, T=temperature, P=pressure)

# #test xsec * path length to transmittance conversion
# nu_hr = np.arange(4100., 4200., 0.00042)
# molecule = "CO"
# Smin = 1.0e-28
# temperature = 210.
# pressure = 100. #Pa = 0.000986923 atm
# xsec = get_xsection(molecule, nu_hr, Smin=Smin, temperature=temperature, pressure=pressure) #cm2/molecule
# # xsec /= 1.0e4 #m2/molecule
# nd = 0.02504e27 * 0.000986923 #number density m-3 at 100 Pa

# for pathl in np.arange(6.0):
#     trans = 10**(-1.49791808045940E+18 * xsec * pathl + 3.38107922285908E-06)
#     plt.plot(nu_hr, trans, label=pathl)
#     print(np.min(trans))
# plt.legend()

def get_molecular_hr(molecule, nu_hr, Smin=0.0, temperature=210.0, pressure=100.0, path_length=100e3): #K, Pa, m

    pressure_atm = pressure * 0.00000986923 #atm

    pressures = {"H2O":113.0*1.0e-6*550, "CO2":550.0, "CO":876.0*1.0e-6*550, "HCl":0.003*1.0e-6*550} #previous: 0.24*1.0e3 pressure in Pa

    if pressure == 0.0:
        pressure = pressures[molecule]
    else:
        pressure *= 1.0e-6*550
    xsec = get_xsection(molecule, nu_hr, Smin=0.0, temperature=temperature, pressure=pressure_atm)
    transmittance = 10**(-1.49791808045940E+18 * xsec * path_length + 3.38107922285908E-06)

    return transmittance


__all__ = [
    "DropBoxHomeResource",
    "DropBoxCollectionResource",
]

from twext.python.log import Logger
from txweb2 import responsecode
from txdav.xml import element as davxml
from txweb2.dav.http import ErrorResponse
from txweb2.dav.resource import DAVResource, TwistedACLInheritable
from txweb2.dav.util import joinURL

from twisted.internet.defer import inlineCallbacks, returnValue

from twistedcaldav.config import config
from twistedcaldav.customxml import calendarserver_namespace

log = Logger()

class DropBoxHomeResource (DAVResource):
    """
    Drop box collection resource.
    """

    def resourceType(self):
        return davxml.ResourceType.dropboxhome  # @UndefinedVariable

    def isCollection(self):
        return True

    def http_PUT(self, request):
        return responsecode.FORBIDDEN

    @inlineCallbacks
    def accessControlList(self, request, *args, **kwargs):
        """
        Override this to give write proxies DAV:write-acl privilege so they can add attachments too.
        """

        acl = (yield super(DropBoxHomeResource, self).accessControlList(request, *args, **kwargs))

        if config.EnableProxyPrincipals:
            owner = (yield self.ownerPrincipal(request))

            newaces = tuple(acl.children)
            newaces += (
                # DAV:write-acl access for this principal's calendar-proxy-write users.
                davxml.ACE(
                    davxml.Principal(davxml.HRef(joinURL(owner.principalURL(), "calendar-proxy-write/"))),
                    davxml.Grant(
                        davxml.Privilege(davxml.WriteACL()),
                    ),
                    davxml.Protected(),
                    TwistedACLInheritable(),
                ),
            )

            returnValue(davxml.ACL(*newaces))

        else:
            returnValue(acl)

class DropBoxCollectionResource (DAVResource):
    """
    Drop box resource.
    """

    def resourceType(self):
        return davxml.ResourceType.dropbox  # @UndefinedVariable

    def isCollection(self):
        return True

    def writeNewACEs(self, newaces):
        """
        Write a new ACL to the resource's property store. We override this for calendar collections
        and force all the ACEs to be inheritable so that all calendar object resources within the
        calendar collection have the same privileges unless explicitly overridden. The same applies
        to drop box collections as we want all resources (attachments) to have the same privileges as
        the drop box collection.

        @param newaces: C{list} of L{ACE} for ACL being set.
        """
        # Add inheritable option to each ACE in the list
        edited_aces = []
        for ace in newaces:
            if TwistedACLInheritable() not in ace.children:
                children = list(ace.children)
                children.append(TwistedACLInheritable())
                edited_aces.append(davxml.ACE(*children))
            else:
                edited_aces.append(ace)

        # Do inherited with possibly modified set of aces
        return super(DropBoxCollectionResource, self).writeNewACEs(edited_aces)

    def http_PUT(self, request):
        return ErrorResponse(
            responsecode.FORBIDDEN,
            (calendarserver_namespace, "valid-drop-box"),
            "Cannot store resources in dropbox",
        )


from typing import Any, Dict, Type, Union

from .. import base
from .model import Model
from .nested import Nested

class Inline(Model, Nested):
    def __init__(
        self,
        props: Dict[
            str,
            Union[Model, base.SchemaAttribute, Type[Model], Type[base.SchemaAttribute]],
        ],
        **kwargs: Any,
    ):
        for key, value in props.items():
            if isinstance(value, type):
                props[key] = value()
        self.__props__ = props
        Nested.__init__(self, model=self, **kwargs)


def validPalindrome(self, s: str) -> bool:
    l, r = 0, len(s) - 1

    while l < r:
        if s[l] != s[r]:
            skipL = s[l + 1 : r + 1]
            skipR = s[l : r]

            if skipL == skipL[::-1] or skipR == skipR[::-1]:
                return True
            else:
                return False
        l += 1
        r -= 1
    return True


# 100% https://app.codility.com/demo/results/trainingDWSUCS-N86/

def solution(A):
    large = []
    small = []
    for index, item in enumerate(A):
        large.append(index + item)
        small.append(index - item)

    large = sorted(large)
    small = sorted(small)

    i = 0
    count = 0
    for index, item in enumerate(large):
        while i < len(large) and item >= small[i]:
            count += i - index
            i += 1

        if count > 10**7:
            return -1

    return count
    '''
    count = 0
    for index, item in enumerate(large):
        for i in range(0, len(large)):
            if item >= small[i]:
                count += i - index
            else:
                pass
        if count > 10**7:
            return -1
    return count
    '''

V. Algorithm/ii. Site/B. Codility/Lessons/Lesson06_04_NumberOfDiscIntersections.py


"""
题目：输入一行字符，分别统计出其中英文字母、空格、数字和其它字符的个数。
"""
def getChrNum(str):
    alph=0
    space=0
    digit=0
    others=0
    for i in str:
        if i.isalpha():
            alph+=1
        elif i.isspace():
            space+=1
        elif i.isdigit():
            digit+=1
        else:
            others+=1
    print("英文字符有%s个，空格有%s个，数字有%s个，其他字符有%s个"%(alph,space,digit,others))

if __name__=="__main__":
    getChrNum("das  451908__")


class Solution:
    def findSubsequences(self, nums: List[int]) -> List[List[int]]:
        # 递归吧
        n = len(nums)
        ans = set()
        def dfs(arr,ind):
            #print(arr,ind)
            if ind == n:
                if len(arr) >= 2:
                    ans.add(tuple(arr))
                return
            if (arr and nums[ind] >= arr[-1]) or not arr:
                dfs(arr + [nums[ind]],ind + 1)
            dfs(arr,ind + 1)

        dfs([],0)
        return list(list(x) for x in ans)


# -*- coding: utf-8 -*-
"""
Created on Mon Nov 12 22:33:31 2018

@author: bpou
"""
#from matplotlib.pyplot import show
import networkx
BJ = 'Beijing'
SZ = 'Shenzhen'
GZ = 'Guangzhou'
WH = 'Wuhan'
HLG = 'Heilongjiang'
NY = 'New York City'
CM = 'Chiangmai'
SG = 'Singapore'

air_route = {
    BJ : {SZ, GZ, WH, HLG, NY},
    GZ : {WH, BJ, CM, SG},
    SZ : {BJ, SG},
    WH : {BJ, GZ},
    HLG : {BJ},
    CM : {GZ},
    NY : {BJ}
}

air_route = networkx.Graph(air_route)

#networkx.draw(air_route, with_labels=True)

def search_destination(graph, start, destination):
    pathes = [[start]]
    visited = set()
    chosen_pathes = []
    while pathes:
        path = pathes.pop(0)
        #print (path)
        frontier = path[-1]
        if frontier in visited:
            continue
        # get new lines

        for city in graph[frontier]:
            new_path = path + [city]
            pathes.append(new_path)
            if city == destination:
               # print (new_path)
                return new_path

        visited.add(frontier)
    return chosen_pathes

def draw_route(cities): return ' ✈ -> '.join(cities)

print(draw_route(search_destination(air_route, SZ, CM)))


import time
from os.path import join

import gymnasium as gym
import numpy as np
import torch
from gymnasium.wrappers import RecordVideo

from main import HIDDEN_LAYER_SIZE, ENV_VERSION, RANDOM_SEED
from net import Net

MODEL_PATH = join("model", "model.pth")

if __name__ == "__main__":
    np.random.seed(RANDOM_SEED)
    env = gym.make(ENV_VERSION, render_mode="human")
    # env.seed(RANDOM_SEED)
    model = Net(env, HIDDEN_LAYER_SIZE)
    model.load_state_dict(torch.load(MODEL_PATH))
    model.eval()
    rewards = 0
    steps = 0
    done = False
    observation = env.reset()[0]
    env = RecordVideo(env=env, video_folder="videos", name_prefix=f"video_{time.time()}")
    while not done:
        action = model.action(observation)
        observation, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        steps += 1
        rewards += reward
        print("Num steps: {} reward {}: ".format(steps, rewards))
    env.close()


""" https://leetcode.com/problems/sum-of-all-odd-length-subarrays/"""

import time
from typing import List

class Solution:
    def sumOddLengthSubarrays(self, arr: List[int]) -> int:
        sum_odd = 0
        n = len(arr)
        for i in range(1, n + 1, 2):
            for j in range(n - i + 1):
                sum_odd += sum(arr[j: j + i])
        return sum_odd

    def sumOddLengthSubarrays2(self, arr: List[int]) -> int:
        res = 0
        n = len(arr)
        for i in range(n):
            res += ((i + 1) * (n - i) + 1) // 2 * arr[i]
        return res

if __name__ == "__main__":
    start_time = time.time()
    print(Solution().sumOddLengthSubarrays([1, 4, 2, 5, 3]))
    print(Solution().sumOddLengthSubarrays([4]))

    print("--- %s seconds ---" % (time.time() - start_time))


import argparse
import os
import sys

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

from torch.autograd import Variable

from generate_input import get_data
from model import ConstrainedModel
import utils

from torch.utils.tensorboard import SummaryWriter
import logging

logging.basicConfig(
    format="%(asctime)s %(levelname)-8s %(message)s",
    level=logging.INFO,
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

from time import time
from collections import defaultdict

def train_nn(
    model,
    train_loader,
    hps,
    model_dir,
    alpha,
    writer,
    beta=None,
):
    # Check Device configuration
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device {device}")

    learning_rate = hps.learning_rate
    n_epochs = hps.n_epochs

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    # Define loss stopping criteria
    ep_stop = 20
    d_loss = []
    conv_loss = 1e-3

    model.to(device)
    logger.info("Training starts!")

    model.train()
    init_wd = []
    for layer in model.layers_weight:
        init_wd += [torch.flatten(layer.data).cpu().numpy()]

    for epoch in range(n_epochs):
        for idx, data in enumerate(train_loader):
            inputs, targets = data
            inputs, targets = inputs.float(), targets.float()

            inputs, targets = (
                Variable(inputs).to(device),
                Variable(targets).to(device),
            )

            # Forward pass
            hidden_states, outputs, _ = model(inputs)

            loss_model = criterion(outputs, targets)

            # Load the constraints to the network
            dist_reg = model.compute_weight_regularizer(norm=hps.norm)
            l_reg = model.compute_lregularizer(norm=hps.norm)

            loss = loss_model + alpha * (
                beta * dist_reg + (1 - beta) * l_reg
            )

            # Backpropagation and optimization
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            writer.add_scalar(
                "training loss",
                loss.item(),
                idx + epoch * len(train_loader),
            )

        if epoch == 0:
            utils.save_checkpoints(
                    {
                        "state_dict": model.state_dict(),
                        "optim_dict": optimizer.state_dict(),
                        "loss": loss,
                        "inputs": inputs,
                        "outputs": outputs,
                        "targets": targets,
                        "hidden_states": hidden_states,
                        "init_wd": init_wd,
                    },
                    checkpoint=os.path.join(
                        model_dir, "checkpoints", "init.pth"
                    ),
                )
        if len(d_loss) >= ep_stop:
            d_loss.pop(0)
        d_loss.append(loss.item())

        if (epoch + 1) % ep_stop == 0:
            acc_avg, acc_max, acc_min, acc_std = utils.evaluate(
                outputs, targets
            )
            accuracies = [acc_avg, acc_max, acc_min, acc_std]
            if (
                np.abs(
                    np.mean(d_loss[: ep_stop // 2])
                    - np.mean(d_loss[ep_stop // 2 :])
                )
                > conv_loss
            ):
                continue
            else:
                logger.info(
                    f"Current epoch is {epoch} and accuracy is {acc_avg}."
                )
                utils.save_checkpoints(
                    {
                        "epoch": epoch,
                        "state_dict": model.state_dict(),
                        "optim_dict": optimizer.state_dict(),
                        "loss": loss,
                        "inputs": inputs,
                        "outputs": outputs,
                        "targets": targets,
                        "hidden_states": hidden_states,
                        "init_wd": init_wd,
                        "accuracy": accuracies,
                    },
                    checkpoint=os.path.join(
                        model_dir, "checkpoints", "last.pth"
                    ),
                )
                print("Checkpoint saved")
                break

def reg_train(hps, model_dir, alpha):
    if isinstance(alpha, (int, float)):
        alpha = [alpha]
    alpha = (
        torch.arange(0.0, 0.009, 0.001)
        if alpha is None
        else np.array(alpha)
    )
    for a in alpha:
        beta = torch.FloatTensor([0.0, 1.0])
        for b in beta:
            logger.info(f"alpha = {a}, beta = {b}")
            # Initialize new dataset and model for each alpha value
            train_loader, test_loader = get_data(hps)
            model = ConstrainedModel(
                hps.n_bits,
                hps.hidden_size,
                hps.n_bits,
                hps.n_spatial_dims,
                hps.norm,
            )

            # Creates directories for each alpha and beta value
            alpha_dir = os.path.join(
                model_dir,
                f"alpha_{round(a.item(),5)}_beta_{round(b.item(),3)}",
            )
            writer = SummaryWriter(log_dir=os.path.join(alpha_dir, "logs"))
            if not os.path.exists(alpha_dir):
                os.makedirs(alpha_dir)
                logger.info(f"Created directory {alpha_dir}")

            # Initialize training
            train_nn(
                model,
                train_loader,
                hps,
                alpha_dir,
                a,
                writer,
                b
            )
    logger.info("Experiments are completed!")

parser = argparse.ArgumentParser(
    description="Process net with distance constraints and L1/L2 regularization."
)
parser.add_argument("--seed", type=int, default=None)
parser.add_argument("--alpha", type=float, default=None)

args = parser.parse_args()

logger.info(f"Passed arguments {args}")

# Initialize the training
manualSeed = args.seed
random.seed(manualSeed)
torch.manual_seed(manualSeed)

torch.cuda.manual_seed(manualSeed)
torch.cuda.manual_seed_all(manualSeed)
torch.backends.cudnn.enabled = False
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

json_path = os.path.join(os.getcwd(), "hps.json")
assert os.path.isfile(
    json_path
), "No json configuration file found at {}".format(json_path)
hps = utils.Params(json_path)

model_dir = os.path.join(
    os.getcwd(), f"trial{hps.trial}", f"seed{manualSeed}trained_model"
)
if not os.path.exists(model_dir):
    os.makedirs(model_dir)
    logger.info(f"Created directory {model_dir}")

reg_train(
    hps,
    model_dir,
    alpha=args.alpha,
)