#!/usr/bin/env python3

# file    :  fuzzman.py
# repo    :  https://github.com/fuzzah/fuzzaide
# author  :  https://github.com/fuzzah
# license :  MIT
# check repository for more information

import os
import sys
import glob
import shutil
import signal
from time import sleep, time
from pprint import pprint

from multiprocessing import cpu_count

from fuzzaide.common import which
from fuzzaide.common.fuzz_stats import is_afl_fuzzer_stats_old, get_afl_stat_name
from .args import get_launch_args
from .running_process import RunningAFLProcess, TimeoutExpired
from .const import *

class FuzzManager:
    """
    Literally GOD-class that needs a ton of rework :(
    Responsible for:
        1. Generating fuzzer commands for different run modes (normal, --builds, --cmd-file)
        2. Running commands
        3. Monitoring and restarting of commands
        4. Statistics counting
        5. Desicion making on termination of fuzzing job

    Limitations:
        1. Only AFL, AFL++ supported
        2. Only interactive output mode supported
    """

    def __init__(self, args):
        self.procs = list()
        self.last_shown_screen_idx = 0
        self.args = args
        self.waited_for_child = False
        self.start_time = int(time())
        self.cores_specified = False
        self.num_from_file = 0

    @staticmethod
    def extract_instance_count(amount):
        """
        Gets number of instances from strings like "5", "10%" or "66.6%"
        """

        count = 0
        perc = amount.endswith("%")
        try:
            if perc:
                count = float(amount.split("%", 1)[0])
            else:
                count = int(amount)
        except ValueError:
            sys.exit(
                "Error in --builds argument: '%s' is not convertible to number of instances (examples: 3, 66.6%%)"
                % (amount,)
            )

        return count, perc

    def extract_complex_mode_params(self):
        """
        Iterate over --builds arguments and extract group name, path, number/percent of cpu cores
        """

        args = self.args

        params = []
        for build_spec in args.builds:
            bspec = build_spec.split(":")  # 0:1:2 -> NAME:PATH:N[%]
            num_spec_parts = len(bspec)
            name = None
            path = None
            count = None
            perc = False
            if num_spec_parts == 1:
                path = bspec[0]
            elif num_spec_parts == 2:
                if "%" in bspec[1] or bspec[1].isnumeric():
                    path = bspec[0]
                    count, perc = self.extract_instance_count(bspec[1])
                else:
                    name = bspec[0]
                    path = bspec[1]
            elif num_spec_parts == 3:
                name = bspec[0]
                path = bspec[1]
                count, perc = self.extract_instance_count(bspec[2])
            else:
                sys.exit(
                    "Error in --builds argument: format of one build is [NAME:]<dir/bin path>[:N[%]] (examples: -h/--help)"
                )

            if path.startswith("~"):
                path = os.path.expanduser(path)
            params.append([name, path, count, perc])

        if args.verbose:
            print("Params of --builds: ")
            pprint(params)

        all_dirs = all(os.path.isdir(p) for _, p, _, _ in params)
        all_bins = all(os.path.isfile(p) for _, p, _, _ in params)

        if not all_dirs and not all_bins:
            sys.exit(
                "Error: --builds should point EITHER to directories OR to binaries"
            )

        # build directories provided -> each one should contain binary with same app name
        if all_dirs:
            for i, (_, dirpath, _, _) in enumerate(params):
                path = os.path.normpath(os.path.join(dirpath, args.program[0]))
                if not os.path.isfile(path):
                    sys.exit(
                        "Error in --builds argument: directory '%s' does not contain '%s' (path checked: '%s')"
                        % (dirpath, args.program[0], path)
                    )
                params[i][1] = path

            return params

        # exact binaries specified (full or partial paths or in PATH)
        for i, (_, binpath, _, _) in enumerate(params):
            if which(binpath) is None:
                sys.exit(
                    "Error in --builds argument: file %s not found so it cannot be tested"
                    % (binpath,)
                )

        return params

    def adjust_complex_mode_params(self, params):
        """
        For complex mode (--builds). Converts percent ratios to number of instances
        and makes sure that each build is used at least once.
        params is a list of 4-item lists: name, path, count/percent, is_percent
        """

        # sum percents as specified by user
        percsum = 0.0
        for _, _, count, is_perc in params:
            if count is not None and is_perc:
                percsum += count

        # normalize percents to 100 if required
        if percsum > 100.0:
            if self.args.verbose:
                print(
                    "Info: sum of percents in --builds is %.2f%% which is not 100%%. Will proportionally adjust it"
                    % (percsum,)
                )
            for i, (_, _, count, is_perc) in enumerate(params):
                if is_perc:
                    params[i][2] = count * 100.0 / percsum
            percsum = 100.0

        # count builds without number or percent of cores specified
        count_none = sum(1 for _, _, count, _ in params if count is None)

        if count_none > 0:
            # set percents for builds without cores specified
            for i, (_, _, count, _) in enumerate(params):
                if count is None:
                    params[i][2] = (100.0 - percsum) / count_none
                    params[i][3] = True

        sum_cores_exact = sum(count for _, _, count, is_perc in params if not is_perc)
        number_of_free_cores = self.args.instances - sum_cores_exact
        if number_of_free_cores < 0 or (number_of_free_cores == 0 and percsum > 0.0):
            sys.exit(
                "Error in --builds argument: not enough processor cores to fit desired configuration"
            )

        if percsum > 0.0:
            for i, (_, _, count, is_perc) in enumerate(params):
                # make 1 core for builds specified as percent
                if is_perc and (
                    count <= 0.0 or number_of_free_cores * params[i][2] / 100.0 < 1
                ):
                    params[i][2] = 1
                    params[i][3] = False

        # final normalization of percents
        percsum = sum(p for _, _, p, is_perc in params if is_perc)
        if percsum != 0.0 and percsum != 100.0:
            for i, (_, _, count, is_perc) in enumerate(params):
                if is_perc:
                    params[i][2] = count * 100.0 / percsum
            percsum = 100.0

        sum_cores_exact = sum(count for _, _, count, is_perc in params if not is_perc)
        number_of_free_cores = self.args.instances - sum_cores_exact
        if number_of_free_cores < 0 or (number_of_free_cores == 0 and percsum > 0.0):
            sys.exit(
                "Error in --builds argument: not enough processor cores to fit desired configuration"
            )

        # convert all the rest percent ratios to number of cores
        if percsum > 0.0:
            for i, (_, _, p, is_perc) in enumerate(params):
                if is_perc:
                    params[i].append(p)
                    params[i][2] = 0
                else:
                    params[i].append(0.0)
                params[i].append(i)  # order

            core_percent = 100.0 / number_of_free_cores
            while True:  # TODO: there must be a better way :(
                params = sorted(params, key=lambda x: -x[4])
                for i, (_, _, _, _, perc, _) in enumerate(params):
                    if perc <= 0.0:
                        continue
                    params[i][4] -= core_percent
                    percsum -= core_percent
                    params[i][2] += 1
                    if percsum <= core_percent:
                        break
                if percsum <= core_percent:
                    break

            number_of_used_cores = sum(c for _, _, c, _, _, _ in params)
            if number_of_used_cores != self.args.instances:
                if self.args.verbose:
                    print(
                        "Math in fuzzman is junky! Fixing error with delta of %d cores"
                        % (number_of_used_cores - self.args.instances,)
                    )
                params = sorted(params, key=lambda x: -x[4])
                params[0][2] -= number_of_used_cores - self.args.instances
                if self.args.verbose:
                    pprint(params)

                number_of_used_cores = sum(c for _, _, c, _, _, _ in params)
                if number_of_used_cores != self.args.instances:
                    sys.exit(
                        "Math in fuzzman is really junky! Error with delta of %d cores! "
                        "Please create an issue with verbose run screenshot.\nFor now you may want to specify different percent/amount of cores"
                        % (number_of_used_cores - self.args.instances,)
                    )

            params = sorted(params, key=lambda x: x[5])  # return original order

        params = list(
            map(lambda x: x[0:3], params)
        )  # leave only name, path and number of cores

        number_of_used_cores = sum(c for _, _, c in params)
        if self.cores_specified and number_of_used_cores != self.args.instances:
            sys.exit(
                "Error in --builds argument: less cores specified in --builds (%d) than in -n (%d)"
                % (number_of_used_cores, self.args.instances)
            )

        if self.args.verbose:
            print(
                "Using %d cores out of total %d available in OS"
                % (number_of_used_cores, cpu_count())
            )

        return params

    def load_custom_cmds(self, path):
        """
        Load custom fuzzer commands from file specified by path.
        Returns list of commands to run instead of fuzzman-generated commands.
        Each element in result list is [worker_name, command]
        """

        cmds = []
        if not os.path.isfile(path):
            sys.exit("Error: file '%s' doesn't exist or it's not a file" % path)

        with open(path, "rt") as f:
            lines = f.readlines()

        for line in lines:
            line = line.strip()
            if len(line) < 1 or line[0] == "#":
                continue

            cmd = line.split(":", 1)
            if len(cmd) != 2:
                sys.exit(
                    "Error: bad command in file %s: %s\nCorrect format:\n  name : command"
                    % (path, line)
                )

            worker = []
            for s in cmd:
                s = s.strip()
                if len(s) < 1:
                    sys.exit(
                        "Error: empty worker name or command in file %s: %s\nCorrect format:\n  name : command"
                        % (path, line)
                    )

                worker.append(s)

            if self.args.verbose:
                print("Loaded custom command: %s" % worker)

            cmds.append(worker)

        # some sanity checks
        if len(cmds) < 1:
            sys.exit("Error: custom commands file doesn't contain any commands to run")

        unique_names = set(name for name, _ in cmds)
        if len(unique_names) < len(cmds):
            sys.exit("Error: custom commands file shouldn't contain duplicate names")

        if not self.args.cmd_file_allow_duplicates:
            unique_cmds = set(cmd for _, cmd in cmds)
            if len(unique_cmds) < len(cmds):
                sys.exit(
                    "Error: custom commands file shouldn't contain duplicate commands (use --cmd-file-allow-duplicates to override)"
                )

        if self.cores_specified and self.args.instances != len(cmds):
            sys.exit(
                "Error: you have specified number of cores = %d but custom commands file %s contains %d commands"
                % (self.args.instances, path, len(cmds))
            )

        return cmds

    def start(self, env=None):
        """
        Start instances either in normal mode, complex mode (--builds) or custom commands mode (--cmd-file)
        """

        if self.args.instances is None:
            self.cores_specified = False
            self.args.instances = cpu_count()
        else:
            self.cores_specified = True

        args = self.args

        if args.instances < 1:
            args.instances = 1

        if which(args.fuzzer_binary) is None:
            sys.exit(
                "File %s not found so it cannot be used as fuzzer" % args.fuzzer_binary
            )

        if not os.path.exists(args.input_dir):
            try:
                os.makedirs(args.input_dir)
            except OSError:
                sys.exit("Can't create input directory %s" % args.input_dir)
        elif not os.path.isdir(args.input_dir):
            sys.exit("Can't use %s as input directory" % args.input_dir)

        if len(glob.glob(os.path.join(args.input_dir, "*"))) < 1:
            path = os.path.join(args.input_dir, "1")
            if args.verbose:
                print("Creating simple input corpus: %s" % path)
            try:
                with open(path, "w") as f:
                    f.write("12345")
            except OSError:
                sys.exit("Wasn't able to create input corpus")

        if args.cleanup and os.path.isdir(args.output_dir):
            print("Removing directory '%s'" % args.output_dir, file=sys.stderr)
            try:
                shutil.rmtree(args.output_dir, ignore_errors=True)
            except shutil.Error:
                sys.exit(
                    "Wasn't able to remove output directory '%s'" % args.output_dir
                )

        if args.cmd_file is not None:
            custom_cmds = self.load_custom_cmds(args.cmd_file)
            for i, (worker_name, cmd) in enumerate(custom_cmds):
                worker_env = os.environ.copy()
                worker_env["AFL_FORCE_UI"] = "1"
                if env is not None:
                    worker_env.update(env)

                print("Starting worker #%d {%s}: %s" % (i + 1, worker_name, cmd))
                self.procs.append(
                    RunningAFLProcess(
                        name=worker_name,
                        groupname="custom",
                        cmd=cmd,
                        env=worker_env,
                        verbose=args.verbose,
                    )
                )
            self.start_time = int(time())
            return

        complex_mode = args.builds is not None and len(args.builds) > 0

        params = []
        used_builds = []
        if complex_mode:
            params = self.extract_complex_mode_params()
            if args.verbose:
                print('"Raw" params:')
                pprint(params)

            params = self.adjust_complex_mode_params(params)
            if args.verbose:
                print("Adjusted params:")
                pprint(params)

            for name, path, num_cores in params:
                used_builds.extend([[name, path]] * num_cores)
        else:  # normal run mode
            if which(args.program[0]) is None:
                sys.exit("File %s not found so it cannot be tested" % args.program[0])
            used_builds = [[None, args.program[0]]] * args.instances

        if args.verbose:
            print("Builds in use:")
            pprint(used_builds)

        # TODO: maybe split this method for basic and complex modes?
        if args.dump_cmd_file:
            print("# Fuzzer commands for use with --cmd-file option of fuzzman")

        for i, (groupname, path) in enumerate(used_builds):
            dictionary = ""

            if i == 0:
                role = "-M"
                worker_name = "m"
                if args.dict:
                    dictionary = " -x " + args.dict

            else:
                role = "-S"
                worker_name = "s"

            power_schedule = " -p explore" if i % 2 == 0 else " -p fast"

            worker_name += str(i + 1)

            if args.no_power_schedules:
                power_schedule = ""

            cmd = (
                args.fuzzer_binary
                + " -i "
                + args.input_dir
                + " -o "
                + args.output_dir
                + " -m "
                + args.memory_limit
                + dictionary
                + " "
                + role
                + " "
                + worker_name
                + power_schedule
            )

            if args.more_args:
                cmd += " " + args.more_args
            cmd += " -- " + path + " " + " ".join(args.program[1:])

            if args.dump_cmd_file:
                worker_env = dict()
            else:
                worker_env = os.environ.copy()
                worker_env["AFL_FORCE_UI"] = "1"

            if env is not None:
                worker_env.update(env)

            cmd = cmd.strip()

            if args.dump_cmd_file:
                wenv = " ".join(k + "=" + v for k, v in worker_env.items())
                if len(wenv) > 0:
                    print("%s : env %s %s" % (worker_name, wenv, cmd))
                else:
                    print("%s : %s" % (worker_name, cmd))
            else:
                print("Starting worker #%d {%s}: %s" % (i + 1, worker_name, cmd))
                self.procs.append(
                    RunningAFLProcess(
                        name=worker_name,
                        groupname=groupname,
                        cmd=cmd,
                        env=worker_env,
                        verbose=args.verbose,
                    )
                )

        if args.dump_cmd_file:
            sys.exit(0)

        self.start_time = int(time())

    def stop(self, grace_sig=signal.SIGINT):
        """
        Stop all fuzzer workers
        """

        if len(self.procs) < 1:
            return

        if self.args.dump_screens:
            print("Dumping status screens")
            self.dump_status_screens()
        else:
            print("Stopping processes")

        self.job_status_check(onlystats=True)

        for proc in self.procs:
            proc.stop(grace_sig=grace_sig)

        term_wait = 1.0
        if self.args.verbose:
            print("Waiting %.1f seconds to check for leftover processes" % (term_wait,))
        sleep(term_wait)

        for proc in self.procs:
            proc.stop(force=True)
        self.procs = []

    def health_check(self):
        """
        Check if fuzzer workers are still running, also print each worker status
        """

        if len(self.procs) < 1:
            return False

        print("Checking status of workers")
        num_ok = sum(1 for proc in self.procs if proc.health_check())

        print("%d/%d workers report OK status" % (num_ok, len(self.procs)))
        return num_ok > 0

    def display_next_status_screen(self, outfile=sys.stdout, static_dump=False):
        """
        Show interactive status screen of one fuzzer for approximately 5 seconds
        """

        if len(self.procs) < 1:
            print("No status screen to show")
            return

        if len(self.procs) == 1:
            self.last_shown_screen_idx = 0

        outbuf = getattr(outfile, "buffer", outfile)

        # helper for drawing workaround on linux with fancy boxes mode
        mqj = b"mqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqj"

        instance = self.procs[self.last_shown_screen_idx]
        if self.is_process_still_running(instance.proc):
            outbuf.write(CURSOR_HIDE)
            if static_dump:
                num_dumps = 1
            else:
                num_dumps = 100

            for _ in range(num_dumps):
                data = instance.get_output(24)
                need_drawing_workaround = not self.args.no_drawing_workaround and len(
                    data
                )

                if need_drawing_workaround:
                    data[0] = data[0].replace(mqj, b"")

                for line in data:
                    outbuf.write(line)

                if need_drawing_workaround:
                    outbuf.write(SET_G1 + bSTG + mqj + bSTOP + cRST + RESET_G1)

                if not static_dump:
                    sleep(0.05)
            outbuf.write(CURSOR_SHOW)
        else:  # process is not running
            if not self.waited_for_child:
                try:
                    instance.proc.wait(3.0)  # wait to prevent zombie-processes
                except TimeoutExpired:
                    pass  # timeout waiting: hanged process?
                else:
                    self.waited_for_child = True

            data = instance.get_output(29)
            for line in data:
                outbuf.write(line)
            if not static_dump:
                sleep(5.0)

        outbuf.write(bSTOP + cRST + RESET_G1 + CURSOR_SHOW)

        if len(self.procs) > 0:
            self.last_shown_screen_idx += 1
            self.last_shown_screen_idx %= len(self.procs)

    @staticmethod
    def is_process_still_running(proc):
        return proc.poll() is None

    def dump_status_screens(self, outfile=sys.stdout):
        """
        Print status screens of all fuzzer instances
        """

        self.last_shown_screen_idx = 0
        outbuf = getattr(outfile, "buffer", outfile)
        for _ in self.procs:
            outbuf.write(
                b"\n" * 40
            )  # messy "workaround" for overlapping status screens (tmux, etc)
            self.display_next_status_screen(outfile=outfile, static_dump=True)

        outbuf.write(b"\n\n")

    def get_fuzzer_stats(self, output_dir, idx, instance):
        """
        Form a dictionary from fuzzer_stats file of given fuzzer instance
        """

        if instance.name is None or len(instance.name) < 1:
            print(
                "Wasn't able to get stats of instance #%d because somehow it has no name"
                % (idx,),
                file=sys.stderr,
            )
            return None

        stats_file_path = os.path.join(output_dir, instance.name, "fuzzer_stats")
        if not os.path.isfile(stats_file_path):
            print(
                "Wasn't able to get stats of instance %s because somehow it has no fuzzer_stats file"
                % (instance.name,),
                file=sys.stderr,
            )
            return None

        try:
            with open(stats_file_path, "rt") as f:
                data = f.readlines()
        except OSError:
            print(
                "Wasn't able to get stats of instance %s because of fail to open '%s'"
                % (instance.name, stats_file_path),
                file=sys.stderr,
            )
            return None

        if data is None:
            print(
                "Wasn't able to get data of instance %s because its fuzzer_stats file '%s' is empty"
                % (instance.name, stats_file_path),
                file=sys.stderr,
            )
            return None

        stats = dict()
        for line in data:
            if ":" not in line:
                continue

            k, v = line.split(":", 1)
            k = k.strip()
            v = v.strip()
            stats[k] = v

        if len(stats) < 1:
            return None

        return stats

    @staticmethod
    def update_stat_timestamp(stats_dict, stat_name, saved_newest_stamp):
        """
        Use this method to update last (newest) path (crash, hang, etc) timestamp.
        Example: newest_path_stamp = update_stat_timestamp(stats, "last_path", newest_path_stamp)
        """

        stamp = stats_dict.get(stat_name)

        if stamp is None:
            return saved_newest_stamp

        try:
            stamp = int(stamp)
        except ValueError:
            return saved_newest_stamp

        if stamp > saved_newest_stamp:
            return stamp

        return saved_newest_stamp

    @staticmethod
    def format_seconds(seconds):
        """
        Returns time in AFL-like format: days, hrs, min, sec
        """

        s = seconds % 60
        m = (seconds // 60) % 60
        h = (seconds // 3600) % 24
        d = seconds // 86400

        if d > 0:
            return "%d days, %d hrs, %d min, %d sec" % (d, h, m, s)
        elif h > 0:
            return "%d hrs, %d min, %d sec" % (h, m, s)
        elif m > 0:
            return "%d min, %d sec" % (m, s)

        return "%d sec" % (s,)

    def job_status_check(self, onlystats=False):
        """
        Enumerate fuzzer_stats files, print stats, return True if stopping required
        """

        output_dir = self.args.output_dir
        if output_dir is None or len(output_dir) < 1:
            return False

        newest_path_stamp = 0
        newest_hang_stamp = 0
        newest_crash_stamp = 0

        sum_execs = 0
        sum_paths = 0
        sum_hangs = 0
        sum_crashes = 0
        sum_restarts = 0

        use_old_style = None

        for idx, instance in enumerate(self.procs, start=1):
            stats = self.get_fuzzer_stats(output_dir, idx, instance)

            if not stats:
                continue

            if use_old_style is None:
                use_old_style = is_afl_fuzzer_stats_old(stats)
                if use_old_style is None:
                    continue

            crashes = int(
                stats.get(get_afl_stat_name("unique_crashes", use_old_style), 0)
            )
            hangs = int(stats.get(get_afl_stat_name("unique_hangs", use_old_style), 0))
            paths_total = int(
                stats.get(get_afl_stat_name("paths_total", use_old_style), 0)
            )
            paths_found = int(
                stats.get(get_afl_stat_name("paths_found", use_old_style), 0)
            )

            sum_restarts += instance.total_restarts
            sum_crashes += crashes
            sum_hangs += hangs
            sum_paths += paths_total
            sum_execs += int(stats.get("execs_done", 0))

            if not onlystats:
                if instance.proc.poll():
                    status = "NOT "
                else:
                    status = ""

                if instance.groupname is not None:
                    print(
                        "Worker %s of group %s is %srunning"
                        % (instance.name, instance.groupname, status)
                    )
                else:
                    print("Worker %s is %srunning" % (instance.name, status))

                print(
                    "\tcrashes: %d, hangs: %d, paths total: %d"
                    % (crashes, hangs, paths_total)
                )
                print(
                    "\tpaths discovered: %d (%.2f%% of total paths)"
                    % (paths_found, 100.0 * paths_found / (paths_total or 1))
                )

            newest_path_stamp = self.update_stat_timestamp(
                stats, get_afl_stat_name("last_path", use_old_style), newest_path_stamp
            )
            newest_hang_stamp = self.update_stat_timestamp(
                stats, get_afl_stat_name("last_hang", use_old_style), newest_hang_stamp
            )
            newest_crash_stamp = self.update_stat_timestamp(
                stats,
                get_afl_stat_name("last_crash", use_old_style),
                newest_crash_stamp,
            )

        print("\nStats of this fuzzing job:")
        job_duration = int(time()) - self.start_time
        print("Duration: %s" % (self.format_seconds(job_duration),))

        if newest_path_stamp == 0:
            if not onlystats:
                print("\nNo more stats to display (yet)")
            return False

        e = float(sum_execs)
        c = ""
        if e >= 1000000000:
            e /= 1000000000
            c = "B"
        elif e >= 1000000:
            e /= 1000000
            c = "M"
        elif e >= 1000:
            e /= 1000
            c = "K"

        if len(c) > 0:
            if c == "B":
                print("   Execs: %.4f%c" % (e, c))
            else:
                print("   Execs: %.2f%c" % (e, c))
        else:
            print("   Execs: %.0f" % (e,))

        now = int(time())

        newest_path_delta = now - newest_path_stamp
        newest_path_fmt = self.format_seconds(newest_path_delta)
        print("   Paths: %d.\tLast new path: %s ago" % (sum_paths, newest_path_fmt))

        if sum_hangs > 0:
            delta = now - newest_hang_stamp
            seconds_fmt = self.format_seconds(delta)
            print("   Hangs: %d.\tLast new hang: %s ago" % (sum_hangs, seconds_fmt))
        else:
            print("   Hangs: 0")

        if sum_crashes > 0:
            delta = now - newest_crash_stamp
            seconds_fmt = self.format_seconds(delta)
            print(" Crashes: %d.\tLast new crash: %s ago" % (sum_crashes, seconds_fmt))
        else:
            print(" Crashes: 0")

        if sum_restarts > 0:
            print("Fuzzer restarts: %d" % (sum_restarts,))

        # now decide if we need to stop
        if (
            self.args.no_paths_stop is not None
            and self.args.no_paths_stop <= newest_path_delta
        ):
            if (
                self.args.minimal_job_duration is not None
                and job_duration < self.args.minimal_job_duration
            ):
                return False
            return True

        return False

def main():
    args = get_launch_args()

    retcode = 7
    fuzzman = FuzzManager(args)

    stdoutbuf = getattr(sys.stdout, "buffer", sys.stdout)

    def handler(_signo, _stack_frame):
        stdoutbuf.write(bSTOP + cRST + RESET_G1 + CURSOR_SHOW)
        print()
        fuzzman.stop()
        sys.exit(0)

    signal.signal(signal.SIGINT, handler)

    fuzzman.start()

    while True:
        stdoutbuf.write(TERM_CLEAR)
        if not fuzzman.health_check():  # this check also prints alive status of workers
            retcode = 1
            break

        sleep(5.0)
        stdoutbuf.write(TERM_CLEAR)
        fuzzman.display_next_status_screen()  # this displays fuzzer output in real time for ~5 seconds

        stdoutbuf.write(TERM_CLEAR)

        # this function displays stats and decides if we need to stop current fuzzing job
        if fuzzman.job_status_check():
            print("STOP CONDITION MET. Stopping current fuzzing job...")
            retcode = 0
            break
        sleep(5.0)

    fuzzman.stop()

    if retcode == 0:
        print("STOP CONDITION MET")

    return retcode

if __name__ == "__main__":
    sys.exit(main())


# Utility, to go through a scan, and see which exact labels were classified incorrectly.

import matplotlib.pyplot as plt
import numpy as np
import sys

import classifier
import cnn
import files
import util

def generateErrorForSingleBrain(scanID):
  data, labelsTrain, labelsTest = files.loadAllInputsUpdated(scanID, classifier.ALL_FEAT, classifier.MORE_FEAT, oneFeat=classifier.ONE_FEAT_NAME)

  trainX, trainY = files.convertToInputs(data, labelsTrain, classifier.PAD, classifier.FLIP_X, classifier.FLIP_Y, classifier.FLIP_Z)
  testX,   testY = files.convertToInputs(data,  labelsTest, classifier.PAD, False, False, False)

  print ("%d train samples, %d test" % (len(trainX), len(testX)))
  _, _, scores, predictions = cnn.runOne(trainX, trainY, testX, testY, scanID, None)
  print ("  Results Train -> Test\n  -------\n" + util.formatScores(scores))

  delta = labelsTest[:, 3] - np.array(predictions)
  errorTest = np.copy(labelsTest.astype(np.float16))
  errorTest[:, 3] = labelsTest[:, 3] - np.array(predictions)
  print ("%f -> %f" % (np.min(errorTest[:, 3]), np.max(errorTest[:, 3])))

  print ("Switched: %d train samples, %d test" % (len(testX), len(trainX)))
  _, _, scores, predictions = cnn.runOne(testX, testY, trainX, trainY, scanID, None)
  print ("  Results Test -> Train\n  -------\n" + util.formatScores(scores))

  delta = labelsTrain[:, 3] - np.array(predictions)
  errorTrain = np.copy(labelsTrain.astype(np.float16))
  errorTrain[:, 3] = labelsTrain[:, 3] - np.array(predictions)
  print ("%f -> %f" % (np.min(errorTrain[:, 3]), np.max(errorTrain[:, 3])))

  errorPath  = "%s/%s/Normal%s-MRA_annotationAll_errors.mat" % (files.BASE_PATH, scanID, scanID)
  print (errorTest.shape)
  print (errorTrain.shape)
  allErrors = np.vstack((errorTrain, errorTest))
  print (allErrors.shape)
  files.writePrediction(errorPath, 'coordTable', allErrors)
  plt.hist(allErrors[:, 3].ravel())
  plt.show()

def main():
  generateErrorForSingleBrain('034')

if __name__ == '__main__':
  classifier.initOptions(sys.argv)
  main()


'''
@inproceedings{DECA:Siggraph2021,
  title={Learning an Animatable Detailed {3D} Face Model from In-The-Wild Images},
  author={Feng, Yao and Feng, Haiwen and Black, Michael J. and Bolkart, Timo},
  journal = {ACM Transactions on Graphics, (Proc. SIGGRAPH)},
  volume = {40},
  number = {8},
  year = {2021},
  url = {https://doi.org/10.1145/3450626.3459936}
}
'''

import os, sys
import numpy as np
from tqdm import tqdm
import torch

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from DECA.deca_model.decalib.deca import DECA
from . import rev_datasets
from DECA.deca_model.decalib.utils.config import cfg as deca_cfg

PATH = "DECA"

class DecaExpGenerator:
    def __init__(self, input_path, device="cuda"):
        self.input_path = input_path
        self.device = device
        self.preprocessed_video_path = os.path.join(PATH, "preprocessed_video")
        self.video_name = os.path.splitext(os.path.basename(input_path))[0]
        self.iscrop = True
        self.detector = "fan"
        self.sample_step = 10
        self.rasterizer_type = "pytorch3d"  # or "standard"

    def generate_expressions_representation(self):
        print("deca start")
        device = self.device

        # load test images
        testdata = rev_datasets.TestData(self.input_path, self.preprocessed_video_path, iscrop=self.iscrop, face_detector=self.detector,
                                     sample_step=self.sample_step)

        # run DECA
        deca_cfg.model.use_tex = False
        deca_cfg.rasterizer_type = self.rasterizer_type
        deca_cfg.model.extract_tex = True
        deca = DECA(config=deca_cfg, device=device)

        exp_array = np.array([])
        frame_counter = 0

        for i in tqdm(range(len(testdata))):
            frame_counter += 1

            name = testdata[i]['imagename']
            images = testdata[i]['image'].to(device)[None, ...]

            with torch.no_grad():
                codedict = deca.encode(images)

                exp_vector = codedict['exp'].cpu().numpy()
                exp_array = np.append(exp_array, exp_vector)

        exp_array = exp_array.reshape(frame_counter, 50)
        return exp_array


from codepack.storages.storable import Storable
from typing import TypeVar, Optional, Union

Code = TypeVar('Code', bound='codepack.code.Code')  # noqa: F821

class Dependency(Storable):
    def __init__(self, code: Optional[Code] = None, id: Optional[str] = None, serial_number: Optional[str] = None,
                 param: Optional[str] = None) -> None:
        super().__init__(id=id, serial_number=serial_number)
        self.code = None
        self.param = None
        self.bind(code)
        self.depend_on(id=id, serial_number=serial_number, param=param)

    def __lshift__(self, sender: Union[Storable, dict]) -> None:
        if isinstance(sender, Storable):
            self.depend_on(id=sender.id, serial_number=sender.serial_number, param=self.param)
        elif isinstance(sender, dict):
            self.depend_on(id=sender['id'], serial_number=sender['serial_number'], param=self.param)
        else:
            raise TypeError(type(sender))
        self.code.add_dependency(self)

    def bind(self, code: Code) -> None:
        self.code = code

    def depend_on(self, id: Optional[str] = None, serial_number: Optional[str] = None,
                  param: Optional[str] = None) -> None:
        self.id = id
        self.serial_number = serial_number
        self.param = param

    def __eq__(self, other: Union['Dependency', dict]) -> bool:
        if isinstance(other, type(self)):
            ret = True
            ret &= (self.serial_number == other.serial_number)
            ret &= (self.id == other.id)
            ret &= (self.param == other.param)
            return ret
        elif isinstance(other, dict):
            ret = True
            ret &= (self.serial_number == other['serial_number'])
            ret &= (self.id == other['id'])
            ret &= (self.param == other['param'])
            return ret
        else:
            return False

    def __str__(self) -> str:
        return '%s(param: %s, id: %s)' % (self.__class__.__name__, self.param, self.id)  # pragma: no cover

    def __repr__(self) -> str:
        return self.__str__()  # pragma: no cover

    def to_dict(self) -> dict:
        return {'id': self.id, 'serial_number': self.serial_number, 'param': self.param}

    @classmethod
    def from_dict(cls, d: dict) -> 'Dependency':
        return cls(id=d['id'], serial_number=d['serial_number'], param=d['param'])


import re
import time
import logging
import numpy as np
import pandas as pd

from selenium import webdriver
from selenium.webdriver.common.by import By
from web_scraper.version import __appName__
from selenium.webdriver.common.action_chains import ActionChains

class WebCollector():
    def __init__(self, web_endpoint: str = '', output_csv_file: str = '', show_graphs: bool = False, driver_path: str = None):
        self.logger = self.setup_logggin()
        self.web_endpoint = web_endpoint
        self.show_graphs = show_graphs
        self.driver_path = driver_path
        self.output_csv_file = output_csv_file
        self.output_fig_path = "../doc/"

        self._main_window_handler = None
        self._secondary_window_handler = None
        self._tertiary_window_handler = None
        self.regex = re.compile('\d{4}S\d')

    # Logger initializier
    def setup_logggin(self) -> logging.Logger:
        log_format = "%(levelname)s %(asctime)s - %(message)s"
        logging.basicConfig(format=log_format, level=logging.INFO)
        logger = logging.getLogger(__appName__)
        return logger

    def load_all_available_dates(self, driver):
        # Find the time button
        time_elem = driver.find_element(by=By.ID, value="TIME")
        time_button = time_elem.find_element(By.TAG_NAME, value='button')

        # We save the main browser window
        self._main_window_handler = driver.current_window_handle
        # Focus on the elements we want to interact
        action = ActionChains(driver)
        action.move_to_element(time_button).perform()
        time_button.click()
        time.sleep(5)

        # Get the secondary window (pop up window)
        while not self._secondary_window_handler:
            for handle in driver.window_handles:
                if handle != self._main_window_handler:
                    self._secondary_window_handler = handle
                    break

        # We change the main window for the pop-up one
        driver.switch_to.window(self._secondary_window_handler)
        # We look for the button to enable all dates
        driver.find_element(by=By.ID, value="checkUncheckAllCheckboxTable").click()
        time.sleep(3)
        # And update the main table which is in the main window
        driver.find_element(by=By.ID, value="updateExtractionButton").click()
        time.sleep(3)
        # Change again to the main window
        driver.switch_to.window(self._main_window_handler)
        time.sleep(3)

    def create_data_table_as_dataframe(self, driver):
        # First we select the preview option so we can iterate over the whole dataset and not only the visible one
        nav_elem = driver.find_element(by=By.ID, value="nav")
        preview_button = nav_elem.find_element(By.CLASS_NAME, value="preview")
        preview_button.click()
        time.sleep(5)

        # Get the third window (pop-up)
        while not self._tertiary_window_handler:
            for handle in driver.window_handles:
                if handle != self._main_window_handler and handle != self._secondary_window_handler:
                    self._tertiary_window_handler = handle
                    break

        driver.switch_to.window(self._tertiary_window_handler)

        table_elem = driver.find_elements(by=By.TAG_NAME, value="table")
        self.logger.info(f"Number of sub-tables {len(table_elem)}")
        self.logger.info("Collecting values from the web site")
        set_columns = []
        table_df = pd.DataFrame()
        # loop over results
        for sub_table in table_elem:
            data = []
            # Remove the "TIME" and "GEO" header and split the table values
            table_values = sub_table.text.split("\n", 1)[1].split("\n")

            # Counting the number of year columns and adding an extra one for the countries
            num_year_cols = len(list(filter(self.regex.match, table_values))) + 1
            # Get the column names
            set_columns = table_values[0:num_year_cols]

            # Iterate over the cost values and jump every num_year_cols columns
            for i in range(num_year_cols, len(table_values), num_year_cols):
                tmp_value = []
                for j in range(0, num_year_cols):
                    tmp_value.append(table_values[i + j])
                data.append(tmp_value)
            # Create a temporal dataframe which is the sub table data
            temporal_df = pd.DataFrame(data=data, columns=set_columns)
            # Merge with the final dataframe
            if table_df.empty:
                table_df = temporal_df
            else:
                table_df = pd.merge(table_df, temporal_df, how="left", on=["GEO"])

        return table_df

    def generate_graph_spain(self, input_csv):
        input_df = pd.read_csv(input_csv, index_col=0)
        year_cols = [x for x in input_df.columns if "GEO" not in x]
        input_df[year_cols] = input_df[year_cols].replace(np.nan, np.float64(0))

        spain_values = input_df.loc[input_df["GEO"] == "Spain"]
        input_df.loc[input_df.index.max() + 1] = input_df[year_cols].mean().replace(np.nan, "European-Mean")
        input_df.iloc[input_df.index.max(), input_df.columns.get_loc('GEO')] = "European-Mean"

        mean_values = input_df.loc[input_df["GEO"] == "European-Mean"]

        spain_values = spain_values.append(mean_values)
        plot_values = spain_values[year_cols].T
        plot_values.columns = ["Spain", "European-Mean"]

        # mean_line = plot_values["European-Mean"].plot(linestyle='-', marker='o')
        plot = plot_values.plot(kind="bar", title="Comparativa de la evolución de precios de la energia Europa vs. España")
        plot.set_xlabel("Año")
        plot.set_ylabel("Precio kWh")
        fig = plot.get_figure()
        fig.savefig(self.output_fig_path + "figure_spain.png")

    def generate_graph_top_expensive(self, input_csv):
        input_df = pd.read_csv(input_csv, index_col=0)
        year_cols = [x for x in input_df.columns if "GEO" not in x]
        input_df[year_cols] = input_df[year_cols].replace(np.nan, np.float64(0))

        countries_to_filter = ["Ireland", "Liechtenstein", "Belgium",
                               "Germany (until 1990 former territory of the FRG)", "Luxembourg", "Italy"]
        input_df = input_df[input_df["GEO"].isin(countries_to_filter)]

        new_col_names = input_df["GEO"]

        plot_values = input_df[year_cols].T
        plot_values.columns = new_col_names

        plot = plot_values.plot(kind="bar", title="Comparativa de la evolución de precios de la energia en Europa para los paises más caros")
        plot.set_xlabel("Año")
        plot.set_ylabel("Precio kWh")
        fig = plot.get_figure()
        fig.savefig(self.output_fig_path + "figure_top_expensives.png")

    def start(self):
        self.logger.info("Starting the web data collector!")
        self.logger.info(f"The website used to do the scraping is: {self.web_endpoint}")

        try:
            self.logger.info("We create the firefox web driver from the path of the executable that we want (if needed)")
            driver = webdriver.Firefox(executable_path=self.driver_path) if self.driver_path else webdriver.Firefox()
            # Maximize the size of the browser to load more JS data
            driver.maximize_window()
            # Selenium don't contain API for request headers
            self.logger.info("Querying the website")
            driver.get(self.web_endpoint)

            self.logger.info("We load all possible years for the data")
            self.load_all_available_dates(driver)

            self.logger.info("We select the preview option to iterate over the table data")
            final_df = self.create_data_table_as_dataframe(driver)

            time.sleep(5)
            self.logger.info("Closing the browser")
            driver.quit()

            self.logger.info("Clean the output dataframe")
            year_cols = [x for x in final_df.columns if "GEO" not in x]
            # Removing strings in price columns and assigning NA for nonexistent ones
            final_df[year_cols] = final_df[year_cols].astype(str).applymap(
                lambda x: re.sub(r'[^0-9^\-\.]+', '', x)).replace('', np.nan).astype('float64')

            self.logger.info(f"Saving the final dataframe in {self.output_csv_file}")
            final_df.to_csv(self.output_csv_file)

            if self.show_graphs:
                self.logger.info("Generating graphs!")
                self.generate_graph_spain(self.output_csv_file)
                self.generate_graph_top_expensive(self.output_csv_file)

        except Exception as e:
            self.logger.warning(f"Exception during the data collection - [{e.args[0]}].")


import uuid
from typing import List

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from schemas import MeetingCreate, Meeting, CommentCreate, Comment
from crud import create_meeting, get_meetings, get_meeting, add_participant_to_meeting, create_comment
from auth.database import User, get_async_session, current_user

router = APIRouter()

@router.post("/meetings/", response_model=Meeting)
async def create_meeting_endpoint(meeting_create: MeetingCreate, user: User = Depends(current_user), db: Session = Depends(get_async_session)):
    meeting = create_meeting(db, meeting_create, user.id)
    if not meeting:
        raise HTTPException(status_code=500, detail="Meeting creation failed")
    return meeting

@router.get("/meetings/", response_model=List[Meeting])
async def get_meetings_endpoint(skip: int = 0, limit: int = 10, db: Session = Depends(get_async_session)):
    meetings = get_meetings(db, skip, limit)
    return meetings

@router.get("/meetings/{meeting_id}", response_model=Meeting)
async def get_single_meeting_endpoint(meeting_id: uuid.UUID, db: Session = Depends(get_async_session)):
    meeting = get_meeting(db, meeting_id)
    if not meeting:
        raise HTTPException(status_code=404, detail="Meeting not found")
    return meeting

@router.post("/meetings/{meeting_id}/add-participant/{user_id}", response_model=Meeting)
async def add_participant_endpoint(meeting_id: uuid.UUID, user_id: uuid.UUID, db: Session = Depends(get_async_session)):
    result = add_participant_to_meeting(db, meeting_id, user_id)
    if not result:
        raise HTTPException(status_code=404, detail="Meeting or user not found")
    return result

@router.post("/meetings/{meeting_id}/comments/", response_model=Comment)
async def create_comment_endpoint(meeting_id: uuid.UUID, comment_create: CommentCreate, user: User = Depends(current_user), db: Session = Depends(get_async_session)):
    comment = create_comment(db, meeting_id, user.id, comment_create)
    if not comment:
        raise HTTPException(status_code=404, detail="Meeting not found")
    return comment


# to-lua.py
# Script created by @Lumaah7 on Twitter
# Lumaah on GitHub
import xml.etree.ElementTree as ET
import sys

def get_animation_prefix(animation_name, prefixes_file):
    with open(prefixes_file, "r") as file:
        valid_prefixes = [line.strip() for line in file]

    for prefix in valid_prefixes:
        if prefix in animation_name.lower():
            return prefix
    return None

def convert_xml_to_lua(xml_file, lua_output_file, prefixes_file, image_path):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    lua_data = "return graphics.newSprite(\n"
    lua_data += f"\tlove.graphics.newImage(graphics.imagePath(\"{image_path}\")),\n"
    lua_data += "\t{\n"

    animation_data = {}
    frame_indices = {}

    for index, sub_texture in enumerate(root.findall(".//SubTexture")):
        name = sub_texture.get("name")
        x = float(sub_texture.get("x"))
        y = float(sub_texture.get("y"))
        width = float(sub_texture.get("width"))
        height = float(sub_texture.get("height"))
        offsetX = float(sub_texture.get("frameX", 0))
        offsetY = float(sub_texture.get("frameY", 0))
        offsetWidth = float(sub_texture.get("frameWidth", width))
        offsetHeight = float(sub_texture.get("frameHeight", height))

        animation_prefix = get_animation_prefix(name, prefixes_file)

        if animation_prefix is None:
            animation_prefix = name.split(" ")[-1]

        if animation_prefix not in frame_indices:
            frame_indices[animation_prefix] = {
                "start": index + 1,
                "stop": index + 1,
                "speed": 12,
                "offsetX": 0,
                "offsetY": 0
            }

        else:
            frame_indices[animation_prefix]["stop"] = index + 1

        lua_data += f"\t\t{{x = {x}, y = {y}, width = {width}, height = {height}, offsetX = {offsetX}, offsetY = {offsetY}, offsetWidth = {offsetWidth}, offsetHeight = {offsetHeight}}}, -- {index+1}: {name}\n"

    lua_data += "\t},\n"
    lua_data += "\t{\n"

    for anim_prefix, indices in frame_indices.items():
        lua_data += f"\t\t[\"{anim_prefix}\"] = {{start = {indices['start']}, stop = {indices['stop']}, speed = {indices['speed']}, offsetX = {indices['offsetX']}, offsetY = {indices['offsetY']}}},\n"

    lua_data += "\t},\n"
    lua_data += "\t\"i\",\n"  # Default animation name (can be changed)
    lua_data += "\tfalse\n"
    lua_data += ")"

    with open(lua_output_file, "w") as lua_file:
        lua_file.write(lua_data)

    print(f"Conversion complete. Output saved to '{lua_output_file}'")

if __name__ == "__main__":
    if len(sys.argv) != 5:
        print("Usage: python script.py input_xml_file output_lua_file animation_prefixes_file image_path")
    else:
        xml_file = sys.argv[1]
        lua_output_file = sys.argv[2] or sys.argv[1].replace(".xml", ".lua")
        prefixes_file = sys.argv[3] or "default_prefixes.txt"
        image_path = sys.argv[4] or "none"
        convert_xml_to_lua(xml_file, lua_output_file, prefixes_file, image_path)


from selenium import webdriver     # 导入webdriver模块

browser_c = webdriver.Chrome()     # 打开谷歌浏览器
browser_c.get("https://www.baidu.com")

'''
print("设置浏览器宽480、高800的显示")
browser_c.set_window_size(480, 800)  #设置窗口高宽
'''

# input = browser_c.find_element_by_id("kw")  # 通过ID定位输入框
input = browser_c.find_element_by_xpath('//*[@id="kw"]')
input.send_keys('python')   # 在输入框输入Python

print(browser_c.current_url)  # 打印当前网页地址
print(browser_c.page_source)  # 打印网页源码

browser_c.close()    #关闭浏览器


import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.utils import save_image

from utils import InstanceNorm2d

class NoiseMappingNetwork(nn.Module):
    def __init__(self, z_dim, hidden_dim, w_dim):
        super(NoiseMappingNetwork, self).__init__()

        # Define the hidden layers using nn.ModuleList
        self.hidden_layers = nn.ModuleList([
            nn.Sequential(
                nn.Linear(hidden_dim, hidden_dim, bias=True),
                nn.LeakyReLU(0.2)
            ) for _ in range(6)
        ])

        # Define the mapping layers using nn.Sequential
        self.mapping = nn.Sequential(
            nn.Linear(z_dim, hidden_dim, bias=True),
            nn.LeakyReLU(0.2)
        )

        for i in range(6):
            self.mapping.add_module(f"hidden_{i}", self.hidden_layers[i])

        self.mapping.add_module("final_layer", nn.Linear(hidden_dim, w_dim, bias=True))

    def forward(self, noise):
        '''
        Function for completing a forward pass of MappingLayers:
        Given an initial noise tensor, returns the intermediate noise tensor.
        '''
        return self.mapping(noise)

class InjectNoise(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.weight = nn.Parameter(
            torch.randn(channels)[None, :, None, None]
        )

    def forward(self, image):
        '''
        Function for completing a forward pass of InjectNoise: Given an image,
        returns the image with random noise added.
        '''
        # Set the appropriate shape for the noise!
        noise_shape = (image.shape[0], 1, image.shape[2], image.shape[3])

        noise = torch.randn(noise_shape, device=image.device) # Creates the random noise
        return image + self.weight * noise # Applies to image after multiplying by the weight for each channel

class AdaIN(nn.Module):
    def __init__(self, channels, w_dim):
        super(AdaIN, self).__init__()

        # Normalize the input per-dimension
        self.instance_norm = InstanceNorm2d(channels)

        self.style_scale_transform = nn.Linear(w_dim, channels)
        self.style_shift_transform = nn.Linear(w_dim, channels)

    def forward(self, image, w):
        '''
        Function for completing a forward pass of AdaIN: Given an image and intermediate noise vector w,
        returns the normalized image that has been scaled and shifted by the style.
        '''
        normalized_image = self.instance_norm(image)
        style_scale = self.style_scale_transform(w)[:, :, None, None]
        style_shift = self.style_shift_transform(w)[:, :, None, None]

        transformed_image = style_scale * normalized_image + style_shift

        return transformed_image

class StyleGANGeneratorBlock(nn.Module):

    def __init__(self, in_chan, out_chan, w_dim, kernel_size=3, use_upsample=True):
        super(StyleGANGeneratorBlock, self).__init__()
        self.use_upsample = use_upsample

        if self.use_upsample:
            self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)

        self.conv1 = nn.Conv2d(in_chan, out_chan, kernel_size, stride=1, padding=1) # Padding is used to maintain the image size
        self.conv2 = nn.Conv2d(out_chan, out_chan, kernel_size, stride=1, padding=1) # Padding is used to maintain the image size

        self.inject_noise1 = InjectNoise(out_chan)
        self.adain1 = AdaIN(out_chan, w_dim)

        self.inject_noise2 = InjectNoise(out_chan)
        self.adain2 = AdaIN(out_chan, w_dim)

        self.activation = nn.LeakyReLU(0.2)

    def forward(self, x, w):
        '''
        Function for completing a forward pass of MicroStyleGANGeneratorBlock: Given an x and w,
        computes a StyleGAN generator block.
        '''
        if self.use_upsample:
            x = self.upsample(x)
        x = self.conv1(x)
        x = self.inject_noise1(x)
        x = self.activation(x)
        x = self.adain1(x, w)

        x = self.conv2(x)
        x = self.inject_noise2(x)
        x = self.activation(x)
        x = self.adain2(x, w)

        return x

class StyleGANGenerator(nn.Module):
    '''
    StyleGAN Generator Class
    Values:
        z_dim: the dimension of the noise vector, a scalar
        map_hidden_dim: the mapping inner dimension, a scalar
        w_dim: the dimension of the intermediate noise vector, a scalar
        channels: a list of channel sizes for each block in the generator
    '''
    def __init__(self, z_dim, map_hidden_dim, w_dim, channels=[512, 256, 128, 64, 32, 16, 8]):
        super(StyleGANGenerator, self).__init__()

        self.num_layers = len(channels) - 1
        self.map = NoiseMappingNetwork(z_dim, map_hidden_dim, w_dim)

        self.initial_block = nn.Parameter(torch.randn(1, channels[0], 4, 4))

        self.blocks = nn.ModuleList([])
        for i in range(self.num_layers):
            in_channels = channels[i]
            out_channels = channels[i + 1]
            block = StyleGANGeneratorBlock(in_channels, out_channels, w_dim, use_upsample=True)  # Fix typo
            self.blocks.append(block)

        self.output_conv_1 = nn.Conv2d(channels[-2], 3, kernel_size=1, stride=1)
        self.output_conv_2 = nn.Conv2d(channels[-1], 3, kernel_size=1, stride=1)
        self.output_tanh = nn.Tanh()

    def upsample_to_match_size(self, smaller_image, bigger_image):
        '''
        Function for upsampling an image to the size of another: Given a two images (smaller and bigger),
        upsamples the first to have the same dimensions as the second.
        Parameters:
            smaller_image: the smaller image to upsample
            bigger_image: the bigger image whose dimensions will be upsampled to
        '''
        return F.interpolate(smaller_image, size=bigger_image.shape[-2:], mode='bilinear')

    def forward(self, noise, return_intermediate=False, alpha=0.2):
        '''
        Function for completing a forward pass of StyleGANGenerator: Given noise,
        computes a StyleGAN iteration.
        '''
        x = self.initial_block
        w = self.map(noise)
        intermediates = []

        for i in range(self.num_layers):
            x = self.blocks[i](x, w)
            if i == self.num_layers - 2:
                x_small = x
                x_small_image = self.output_tanh(self.output_conv_1(x_small))
                intermediates.append(x_small_image)
            if i == self.num_layers - 1:
                x_big = x
                x_big_image = self.output_tanh(self.output_conv_2(x_big))
                intermediates.append(x_big_image)

        x_small_upsample = self.upsample_to_match_size(intermediates[0], intermediates[1])

        interpolation = alpha * x_big_image + (1 - alpha) * x_small_upsample

        if return_intermediate:
            return interpolation, x_small_upsample, intermediates
        else:
            return interpolation

if __name__ == "__main__":
    # Set the required parameters
    z_dim = 512
    map_hidden_dim = 1024
    w_dim = 512
    channels = [512, 256, 128, 64, 32, 16, 8]

    # Create an instance of the generator
    generator = StyleGANGenerator(z_dim, map_hidden_dim, w_dim, channels)

    # Generate a random noise vector
    batch_size = 4
    noise = torch.randn(batch_size, z_dim)

    # Generate images using the generator blocks
    images = generator(noise)
    print(images.size())

    # Save the generated images
    save_image(images, 'generated_images.png', nrow=2, normalize=True)


import numpy as np
from scipy.ndimage import map_coordinates

def rectify_orders(image_data, trace, half_window=10, debug=False, nullify_mapped_values=True):
    """
    :param image_data: xwavecal.images.Image.data
    :param trace: Trace object
    :param half_window: half of the size of the extraction window about the center of the traces. E.g. 10 means 21 pixel
                        window will be extracted.
    :param debug: boolean for returning an internal copy of the image which is used in extract.
    :param nullify_mapped_values: bool.
                                  True if the flux from a pixel can only be used in one order's spectrum.
                                  False if a pixel's flux can be double (or triple etc..) counted.
    :return rectified_orders: Dictionary
    Dictionary where the keys are the trace id's from trace.get_id(),
    where rectified_2d_spectrum[trace_id]['val'] is a 2d float array (flux for the trace_id order).
    If half extraction window was 10, then rectified_2d_spectrum[trace_id]['val']
    is 21 rows by 4096 columns (for a 4096 pixel wide image). One would sum this 2d
    spectrum along-columns to get a box extracted spectrum.
    """
    rectified_orders = {}
    num_orders = trace.num_traces_found()
    image_copy = np.copy(image_data)
    x_coordinates, y_coordinates = np.meshgrid(np.arange(image_data.shape[1]), np.arange(image_data.shape[0]))
    image_coordinates = {'x': x_coordinates, 'y': y_coordinates}
    for index in np.arange(num_orders):
        rectified_orders[trace.get_id(index)], image_copy = rectify_order(image_copy, image_coordinates,
                                                                          trace.get_centers(index),
                                                                          half_window=half_window,
                                                                          nullify_mapped_values=nullify_mapped_values)
    if debug:
        return rectified_orders, image_copy
    else:
        return rectified_orders

def rectify_order(image_data, image_coordinates, single_order_centers, half_window=10, nullify_mapped_values=True):
    """
    :param image_data: xwavecal.images.Image.data
    :param image_coordinates: dictionary with x and y coordinates for each pixel.
    :param single_order_centers: the y centers of the center of the rectification window
    :param half_window: the half width of the window about which one is rectifying the trace/order
    :param nullify_mapped_values: True if the flux from a pixel can only be used in one order's spectrum.
                                  False if a pixel's flux can be double (or triple etc..) counted.
    :return: An array of 2*half_window + 1 rows by width(image_data) about the order. The center of the order
             is right at the index of half_window, e.g. rectified_order[half_window] gives the flux down the
             center of the order.
    """
    spectrum_shape = (2*half_window + 1, image_data.shape[1])
    rectified_order = {'val': np.zeros(spectrum_shape),
                       'y': np.zeros(spectrum_shape),
                       'x': np.zeros(spectrum_shape)}
    x_coords = np.arange(image_data.shape[1])
    for offset, row in zip(np.arange(-half_window, half_window + 1), np.arange(spectrum_shape[0])):
        mapped_y_values = map_coordinates(image_coordinates['y'], [single_order_centers + offset, x_coords],
                                          order=0, mode='constant', cval=0, prefilter=False)
        rectified_order['val'][row] = image_data[(mapped_y_values, x_coords)]
        rectified_order['y'][row] = mapped_y_values
        rectified_order['x'][row] = x_coords
        if nullify_mapped_values:
            image_data[(mapped_y_values, x_coords)] = 0
    return rectified_order, image_data


import os
os.environ['GDAL_SKIP'] = 'DODS'
import logging
import isce
root_logger = logging.getLogger()
root_logger.setLevel('WARNING')
from carst import SingleRaster, ConfParams
from carst.libft import ampcor_task, writeout_ampcor_task
from carst.libxyz import AmpcoroffFile

import pandas as pd
import ipyleaflet as ilfl
import ipywidgets as iwg
from geostacks import SpatialIndexLS8, SpatialIndexITSLIVE
import xarray as xr
import rasterio
import fsspec
from datetime import datetime

import subprocess

# CARST monkey patch (to temporarily resolve .tiff vs .tif issue)

def Unify(self, params):
    """ Use gdalwarp to clip, reproject, and resample a DEM using a given set of params (which can 'unify' all DEMs). """

    print('Calling Gdalwarp...')
    fpath_warped = self.fpath[:-5] + '_warped' + self.fpath[-5:]
    if 'output_dir' in params:
        newpath = '/'.join([params['output_dir'], fpath_warped.split('/')[-1]])
    else:
        # for pixel tracking (temporarily)
        newfolder = 'test_folder'
        if not os.path.exists(newfolder):
            os.makedirs(newfolder)
        newpath = '/'.join([newfolder, fpath_warped.split('/')[-1]])
        newpath = newpath.replace(newpath.split('.')[-1], 'img')
    gdalwarp_cmd = 'gdalwarp -t_srs ' + params['t_srs'] + ' -tr ' + params['tr'] + ' -te ' + params['te']
    if 'of' in params:
        gdalwarp_cmd += ' -of ' + params['of']
    if 'ot' in params:
        gdalwarp_cmd += ' -ot ' + params['ot']
    gdalwarp_cmd += ' -overwrite ' + self.fpath + ' ' + newpath
    print(gdalwarp_cmd)
    retcode = subprocess.call(gdalwarp_cmd, shell=True)
    if retcode != 0:
        print('Gdalwarp failed. Please check if all the input parameters are properly set.')
        sys.exit(retcode)
    self.fpath = newpath

SingleRaster.Unify = Unify

# Adapted from Alice's AGU talk

class eztrack_ui():

    def __init__(self, zoom=4, query_pt=[-50., 69.], spatial_index=None):

        self.zoom = zoom
        self.query_pt = query_pt
        self.mainmap = None
        self.marker = None
        self.idxs = None
        self.ui_title = None
        self.prlist = []
        self.menuleft = None     # path / row menu
        self.pr_selection = None
        self.kernelselection = None
        self.datesearch_btn = None
        self.scenelist = []
        self.menuright = None    # scene menu
        self.bandselection = None
        self.runft_btn = None
        self.map_polygon = None
        self.spatial_index = spatial_index
        self.output = iwg.Output()   # print message output
        self.results = None
        self.ft_params = None     # feature tracking parameters
        self.sld1 = None         # param slider #1
        self.sld2 = None         # param slider #2
        self.sld3 = None         # param slider #3

    def init_panelleft(self):
        self.ui_title = iwg.HTML("
Drag the marker to your region of interest
")
        self.idxs = self.spatial_index.query_pathrow(self.query_pt)
        self.prlist = [('{:03d}/{:03d}'.format(self.spatial_index.footprint.loc[i, 'path'], self.spatial_index.footprint.loc[i, 'row']) ,i) for i in self.idxs]
        self.menuleft = iwg.Select(options=self.prlist, description='Path/Row:', rows=15)
        self.kernelselection = iwg.RadioButtons(options=['ITS_LIVE (online ready)', 'CARST'], value='ITS_LIVE (online ready)', description='Data / Kernel:')
        self.datesearch_btn = iwg.Button(description='Search for dates')

    def init_panelright(self):
        self.menuright = iwg.SelectMultiple(options=self.scenelist, description='Data entries:', rows=20)
        self.bandselection = iwg.RadioButtons(options=['B4', 'B8'], value='B8', description='Band (LS8):')
        self.runft_btn = iwg.Button(description='Get data / Start feature tracking')

    def init_map(self):
        self.mainmap = ilfl.Map(basemap=ilfl.basemaps.Gaode.Satellite, center=[self.query_pt[-1], self.query_pt[0]], zoom=self.zoom)
        self.marker = ilfl.Marker(location=[self.query_pt[-1], self.query_pt[0]], draggable=True)
        self.mainmap.add_layer(self.marker)
        self.pr_selection = self.idxs[0]
        self.scene_list = pd.DataFrame(columns=('prefix', 'time', 'tier'))
        self.map_polygon = ilfl.WKTLayer(wkt_string=self.spatial_index.footprint.loc[self.pr_selection].geometry.wkt)
        self.mainmap.add_layer(self.map_polygon)

    def gen_ui(self, spatial_index=None, ft_params=None):
        if self.spatial_index is None:
            self.spatial_index = spatial_index

        if ft_params is None:
            self.init_ft_params()
        else:
            self.ft_params = ft_params

        self.init_panelleft()
        self.init_panelright()
        self.init_map()

        self.marker.observe(self._on_location_changed, 'location')
        self.menuleft.observe(self._on_menuleft_selection_changed, names='value')
        self.datesearch_btn.on_click(self._on_searchbutton_clicked)
        self.runft_btn.on_click(self._on_ftbutton_clicked)

        leftside = iwg.VBox([self.ui_title, self.menuleft, self.kernelselection, self.datesearch_btn])
        leftside.layout.align_items = 'center'
        rightside = iwg.VBox([self.menuright, self.bandselection, self.runft_btn])
        rightside.layout.align_items = 'center'
        return iwg.AppLayout(left_sidebar=leftside, center=self.mainmap, right_sidebar=rightside)

    # ==== leftmenu update when map marker loc changes

    def _on_location_changed(self, event):
        # global query_pt, idx
        self.query_pt = [self.marker.location[-1], self.marker.location[0]]
        self.idxs = self.spatial_index.query_pathrow(self.query_pt)
        self.prlist = [('{:03d}/{:03d}'.format(self.spatial_index.footprint.loc[i, 'path'], self.spatial_index.footprint.loc[i, 'row']) ,i) for i in self.idxs]
        self.menuleft.options = self.prlist

    # ==== map polygon update when leftmenu selection changes

    def _on_menuleft_selection_changed(self, change):
        # global pr_selection
        self.pr_selection = change['new']
        self.map_polygon.wkt_string=self.spatial_index.footprint.loc[self.pr_selection].geometry.wkt

    # ==== search button click callback

    def _on_searchbutton_clicked(self, event):
        # global pr_scene_list
        with self.output:
            if self.kernelselection.value == 'CARST':
                s3_prefix, self.scenelist = self.spatial_index.search_s3(self.pr_selection)
                # print(s3_prefix)
                self.menuright.options = [(record['time'], idx) for idx, record in self.scenelist.loc[self.scenelist['tier'] == 'T1'].iterrows()]
            elif self.kernelselection.value == 'ITS_LIVE (online ready)':
                # query_pt = [mker.location[-1], mker.location[0]]
                polygon_coords = SpatialIndexITSLIVE.get_minimal_bbox(self.query_pt)
                params = {'polygon': polygon_coords, 'percent_valid_pixels': 1, 'start': '2015-01-01', 'end': '2020-01-01'}
                urls = SpatialIndexITSLIVE.get_granule_urls(params)
                self.scenelist = SpatialIndexITSLIVE.parse_urls(urls)
                pr_dict_key = '{:03d}/{:03d}'.format(self.spatial_index.footprint.loc[self.pr_selection, 'path'], self.spatial_index.footprint.loc[self.pr_selection, 'row'])
                self.menuright.options = [(i['entrystr'], i['url']) for i in self.scenelist[pr_dict_key]]

    # ==== feature tracking button click callback

    # file1_url  = None
    # file1_date = None
    # file2_url  = None
    # file2_date = None

    def _on_ftbutton_clicked(self, ft):
        # global file1_url, file1_date, file2_url, file2_date
        with self.output:
            if self.kernelselection.value == 'CARST':
                selected_list = self.scenelist.loc[list(self.menuright.value)]
                selected_list_prefix = selected_list['prefix'].tolist()
                selected_list_time = selected_list['time'].tolist()
                file1_url = 'https://landsat-pds.s3.amazonaws.com/' + selected_list_prefix[0] + os.path.basename(selected_list_prefix[0][:-1]) + '_' + self.bandselection.value + '.TIF'
                file2_url = 'https://landsat-pds.s3.amazonaws.com/' + selected_list_prefix[1] + os.path.basename(selected_list_prefix[1][:-1]) + '_' + self.bandselection.value + '.TIF'
                file1_date = selected_list_time[0].strftime('%Y-%m-%d')
                file2_date = selected_list_time[1].strftime('%Y-%m-%d')
                print(file1_url, file1_date, file2_url, file2_date)
                carst_featuretrack(file1_url, file1_date, file2_url, file2_date, params=self.ft_params)
                self.results = rasterio.open(selected_list_time[0].strftime('%Y%m%d') + '-' + selected_list_time[1].strftime('%Y%m%d') + '_velo-raw_mag.tif')
            elif self.kernelselection.value == 'ITS_LIVE (online ready)':
                print(self.menuright.value)
                with fsspec.open(self.menuright.value[0]) as fobj:
                    self.results = xr.open_dataset(fobj)

    # ==== Initialize feature tracking parameters

    def init_ft_params(self):
        self.ft_params = FTParams()

    def show_ft_params(self):
        self.sld1 = iwg.IntSlider(value=self.ft_params.pxsettings['refwindow_x'], step=16, min=16, max=256, description='reference window size')
        self.sld2 = iwg.IntSlider(value=self.ft_params.pxsettings['searchwindow_x'], step=1, min=16, max=128, description='search window size')
        self.sld3 = iwg.IntSlider(value=self.ft_params.pxsettings['skip_across'], step=1, min=1, max=256, description='skip size')
        self.sld1.observe(self._on_param_change_sld1, names='value')
        self.sld2.observe(self._on_param_change_sld2, names='value')
        self.sld3.observe(self._on_param_change_sld3, names='value')
        return iwg.VBox([self.sld1, self.sld2, self.sld3])

    def _on_param_change_sld1(self, change):
        self.ft_params.pxsettings['refwindow_x'] = self.sld1.value
        self.ft_params.pxsettings['refwindow_y'] = self.sld1.value

    def _on_param_change_sld2(self, change):
        self.ft_params.pxsettings['searchwindow_x'] = self.sld2.value
        self.ft_params.pxsettings['searchwindow_y'] = self.sld2.value

    def _on_param_change_sld3(self, change):
        self.ft_params.pxsettings['skip_across'] = self.sld3.value
        self.ft_params.pxsettings['skip_down'] = self.sld3.value

class FTParams:

    def __init__(self):
        self.imagepair     = {'image1': None,
                              'image1_date': None,
                              'image2': None,
                              'image2_date': None}
        self.pxsettings    = {'refwindow_x': 64,
                              'refwindow_y': 64,
                              'searchwindow_x': 20,
                              'searchwindow_y': 20,
                              'skip_across': 64,
                              'skip_down': 64,
                              'oversampling': 16,
                              'threads': 8,
                              'gaussian_hp': False,
                              'gaussian_hp_sigma': 3.0}
        self.outputcontrol = {'datepair_prefix': True,
                              'output_folder': '.'}
        self.rawoutput     = {'if_generate_ampofftxt': False,
                              'if_generate_xyztext': False,
                              'label_ampcor': 'ampoff',
                              'label_geotiff': 'velo-raw'}

    def VerifyParams(self):

        """
        Verify params and modify them to proper types.
        """

        if hasattr(self, 'pxsettings'):
            if 'size_across' not in self.pxsettings:
                self.pxsettings['size_across'] = None
            if 'size_down' not in self.pxsettings:
                self.pxsettings['size_down'] = None

        if hasattr(self, 'outputcontrol'):
            if 'datepair_prefix' in self.outputcontrol:
                if self.outputcontrol['datepair_prefix'] in ['false', 'f', 'no', 'n', '0']:
                    self.outputcontrol['if_generate_xyztext'] = False
                else:
                    self.outputcontrol['datepair_prefix'] = bool(int(self.outputcontrol['datepair_prefix']))
                    atime = datetime.strptime(self.imagepair['image1_date'], '%Y-%m-%d')
                    btime = datetime.strptime(self.imagepair['image2_date'], '%Y-%m-%d')
                    self.outputcontrol['label_datepair'] = atime.strftime('%Y%m%d') + '-' + btime.strftime('%Y%m%d' + '_')
        if hasattr(self, 'rawoutput'):
            if 'label_ampcor' in self.rawoutput:
                if self.outputcontrol['datepair_prefix'] not in ['false', 'f', 'no', 'n', '0']:
                    self.rawoutput['label_ampcor'] = os.path.join(self.outputcontrol['output_folder'], self.outputcontrol['label_datepair'] + self.rawoutput['label_ampcor'])
                else:
                    self.rawoutput['label_ampcor'] = os.path.join(self.outputcontrol['output_folder'], self.rawoutput['label_ampcor'])
            if 'label_geotiff' in self.rawoutput:
                if self.outputcontrol['datepair_prefix'] not in ['false', 'f', 'no', 'n', '0']:
                    self.rawoutput['label_geotiff'] = os.path.join(self.outputcontrol['output_folder'], self.outputcontrol['label_datepair'] + self.rawoutput['label_geotiff'])
                else:
                    self.rawoutput['label_geotiff'] = os.path.join(self.outputcontrol['output_folder'], self.rawoutput['label_geotiff'])

# CARST feature tracking workflow

def carst_featuretrack(file1_url, file1_date, file2_url, file2_date, params=None, inipath='param.ini'):
    if params is not None:
        ini = params
    else:
        ini = ConfParams(inipath)
        ini.ReadParam()
    ini.imagepair['image1'] = file1_url
    ini.imagepair['image2'] = file2_url
    ini.imagepair['image1_date'] = file1_date
    ini.imagepair['image2_date'] = file2_date
    ini.VerifyParams()
    a = SingleRaster(file1_url, date=file1_date)
    b = SingleRaster(file2_url, date=file2_date)
    # ========== Unifying LS8 image extent ==========
    te_a = a.GetExtent()
    tr_a = (a.GetXRes(), -a.GetYRes())
    t_srs_a = a.GetProj4()
    params_unify = {'output_dir': '.', 't_srs': '"' + t_srs_a + '"', 'tr': '{} {}'.format(*tr_a), 'te': '{} {} {} {}'.format(te_a[0], te_a[3], te_a[2], te_a[1])}
    b.Unify(params_unify)
    # ===============================================
    if ini.pxsettings['gaussian_hp']:
        a.GaussianHighPass(sigma=ini.pxsettings['gaussian_hp_sigma'])
        b.GaussianHighPass(sigma=ini.pxsettings['gaussian_hp_sigma'])
    a.AmpcorPrep()
    b.AmpcorPrep()
    task = ampcor_task([a, b], ini)
    writeout_ampcor_task(task, ini)
    ampoff = AmpcoroffFile(ini.rawoutput['label_ampcor'] + '.p')
    ampoff.Load()
    ampoff.SetIni(ini)
    ampoff.FillwithNAN()   # fill holes with nan
    ampoff.Ampcoroff2Velo()
    ampoff.Velo2XYV(generate_xyztext=ini.rawoutput['if_generate_xyztext'])
    ampoff.XYV2Raster()


"""Main app/routing for Twitter"""

from flask import Flask, render_template
from models import DB, User, Tweet
from twitter import *

# function to intialize the flask instance

# create the application
def create_app():
    """Create and configuring an instance of the Flask application"""

    # location of the Flask application and look in the current directory
    app = Flask(__name__)

    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///db.sqlite3'
    app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
    # with app.app_context():
    DB.init_app(app)

    # DB.init_app(app)  # initialize the DB with this application, DB is an instance of SQLAlchemy class

    # add a decorator
    @app.route('/')
    def root():
        # DB.drop_all()
        # DB.create_all()
        # app_user = User(id=1, name='app_user')
        # DB.session.add(app_user)
        # DB.session.commit()

        # return render_template('base.html',
        #                        title='Home',
        #                        users=User.query.all())
        return "text"
        # PS muse: when you are using Title = Home you are using Jinja2

    return app

#
# if __name__ == '__main__':
#     create_app().run()


from rest_framework import serializers
from django.db.models import Q, QuerySet

from certificates.models import StudentCertificate
from courses.models import UserCourse
from users.serializers import UserMinSerializer
from .models import Program, Field, UserBuyProgram, StudentProgram, UserViewProgram
from courses.serializers import CourseSerializer, CourseSearchSerializer, CourseMinSerializer, CourseProcessSerializer, \
    UserCourseSerializer

class ProgramSerializer(serializers.ModelSerializer):
    id = serializers.IntegerField(read_only=True)
    field = serializers.PrimaryKeyRelatedField(many=False, read_only=True)
    created_by = serializers.PrimaryKeyRelatedField(many=False, read_only=True)
    program_course = CourseSerializer(many=True, read_only=True)

    class Meta:
        model = Program
        fields = [
            'id', 'name',  'icon', 'slug', 'program_course',
            'status', 'field', 'created_date', 'discount_percentage',
            'short_description', 'full_description', 'created_date', 'created_by',
            'created_by_name', 'modified_date'
        ]
        read_only_fields = ('created_date', 'modified_date',
                            'created_by', 'created_by_name')

class ProgramDetailSerializer(serializers.ModelSerializer):
    id = serializers.IntegerField(read_only=True)
    field = serializers.StringRelatedField(read_only=True)
    program_course = CourseSearchSerializer(many=True, read_only=True)
    courses_count = serializers.IntegerField(read_only=True)
    is_my_program = serializers.SerializerMethodField()
    price = serializers.SerializerMethodField()
    bought_courses = serializers.SerializerMethodField(read_only=True)
    completed_courses = serializers.SerializerMethodField(read_only=True)

    class Meta:
        model = Program
        fields = [
            'id', 'name', 'icon', 'slug', 'discount_percentage', 'price',
            'program_course', 'benefits', 'pre_requisites', 'completed_courses',
            'courses_count', 'status', 'field', 'bought_courses',
            'short_description', 'full_description', 'is_my_program'
        ]

    def get_is_my_program(self, obj):
        user = self.context.get('user')
        return obj.user_buy.filter(pk=user.id).count() > 0

    def get_price(self, obj):
        price = 0
        user = self.context.get('user')
        program_courses = obj.program_course.all()
        for course in program_courses:
            if course.check_is_bought(user) is not True:
                price += int(course.get_price())
        # if obj.discount:
        #     price -= int(obj.discount)
        if obj.discount_percentage:
            price *= (100-obj.discount_percentage)/100
        return int(price)

    def get_bought_courses(self, obj):
        user = self.context.get('user')
        if user and not user.is_anonymous:
            program_courses = obj.program_course.all()
            queryset = [course for course in program_courses if course.check_is_bought(student=user)]
            return CourseMinSerializer(instance=queryset, many=True).data
        return []

    def get_completed_courses(self, obj):
        user = self.context.get('user')
        if user and not user.is_anonymous:
            program_courses = obj.program_course.all()
            queryset = [course for course in program_courses if course.check_is_completed(student=user)]
            return CourseMinSerializer(instance=queryset, many=True).data
        return []

class ProgramSearchSerializer(serializers.ModelSerializer):

    id = serializers.IntegerField(read_only=True)
    field = serializers.StringRelatedField(read_only=True)
    courses_count = serializers.IntegerField(read_only=True)
    program_course = CourseSearchSerializer(many=True, read_only=True)
    status = serializers.SerializerMethodField(read_only=True)

    class Meta:
        model = Program
        fields = [
            'id', 'name', 'icon', 'slug', 'discount_percentage', 'status',
            'courses_count', 'status', 'field', 'program_course'
        ]

    def get_status(self, obj):
        request = self.context.get('request', None)
        if request and request.user is not None and not request.user.is_anonymous:
            try:
                student_program = StudentProgram.objects.get(student=request.user, program_id=obj.id)
                return student_program.status
            except StudentProgram.DoesNotExist:
                return None
        return None

class ProgramMinSerializer(serializers.ModelSerializer):
    id = serializers.IntegerField(read_only=True)
    field = serializers.StringRelatedField(read_only=True)
    courses_count = serializers.IntegerField(read_only=True)
    program_course = CourseMinSerializer(many=True, read_only=True)
    bought_date = serializers.SerializerMethodField(read_only=True)

    class Meta:
        model = Program
        fields = [
            'id', 'name', 'icon', 'slug', 'bought_date', 'discount_percentage',
            'courses_count', 'status', 'field', 'program_course'
        ]

    def get_bought_date(self, obj):
        user = self.context.get('user')
        if user is not None:
            try:
                instance = UserBuyProgram.objects.get(user_id=user.id, program_id=obj.id)
            except UserBuyProgram.DoesNotExist:
                return None
            return instance.bought_date
        return None

class ProgramProcessSerializer(serializers.ModelSerializer):

    id = serializers.IntegerField(read_only=True)
    field = serializers.StringRelatedField(read_only=True)
    courses_count = serializers.IntegerField(read_only=True)
    program_course = CourseProcessSerializer(many=True, read_only=True)
    student_program = serializers.SerializerMethodField(read_only=True)
    student_course = serializers.SerializerMethodField(read_only=True)

    class Meta:
        model = Program
        fields = [
            'id', 'name', 'icon', 'slug', 'discount_percentage', 'student_program', 'student_course',
            'courses_count', 'status', 'field', 'program_course',
        ]

    def get_student_program(self, obj):
        user = self.context.get('user')
        request = self.context.get('request')
        context = {"user": user, "request": request}
        try:
            instance = StudentProgram.objects.get(student_id=user.id, program_id=obj.id)
            return StudentProgramSerializer(instance=instance, context=context).data
        except StudentProgram.DoesNotExist:
            return None

    def get_student_course(self, obj):
        user = self.context.get('user')
        queryset = list()
        program_courses = obj.program_course.all()
        for course in program_courses:
            try:
                instance = UserCourse.objects.get(Q(user_id=user) & Q(course_id=course.id))
                queryset.append(instance)
            except UserCourse.DoesNotExist:
                pass
        return UserCourseSerializer(instance=queryset, many=True).data

class ProgramDataSerializer(serializers.ModelSerializer):
    view_count = serializers.SerializerMethodField(read_only=True)
    courses_count = serializers.IntegerField(read_only=True)

    class Meta:
        model = Program
        fields = ['id', 'created_date', 'name', 'view_count',
                  'courses_count', 'field', 'status', 'discount_percentage', 'discount_percentage']

    @staticmethod
    def get_view_count(obj):
        return UserViewProgram.objects.filter(program_id=obj.id).count()

class FieldSerializer(serializers.ModelSerializer):
    id = serializers.IntegerField(read_only=True)
    field_programs = ProgramSearchSerializer(many=True, read_only=True)
    field_courses = CourseSearchSerializer(many=True, read_only=True)
    field = serializers.StringRelatedField(read_only=True)

    class Meta:
        model = Field
        fields = [
            'id', 'name', 'icon', 'description', 'field_programs', 'field_courses', 'field', 'created_date'
        ]
        read_only_fields = ('created_date',)

class FieldMinSerializer(serializers.ModelSerializer):
    id = serializers.IntegerField(read_only=True)
    slug = serializers.SlugField(read_only=True)

    class Meta:
        model = Field
        fields = [
            'id', 'name',  'slug', 'icon', 'description'
        ]

class UserBuyProgramSerializer(serializers.ModelSerializer):
    user = serializers.StringRelatedField(read_only=True, required=False)
    program = serializers.StringRelatedField()

    class Meta:
        model = UserBuyProgram
        fields = [
            'id', 'user', 'program', 'bought_date', 'money'
        ]

class StudentProgramSerializer(serializers.ModelSerializer):
    id = serializers.IntegerField(read_only=True)
    student = UserMinSerializer(required=False)
    program = ProgramMinSerializer(required=False)
    course_num = serializers.SerializerMethodField(read_only=True)
    # file = serializers.SerializerMethodField(read_only=True)

    class Meta:
        model = StudentProgram
        fields = [
            'id', 'file', 'student', 'program', 'status',
            'started_date', 'completed_date', 'received_certificate', 'course_num'
        ]

    @staticmethod
    def get_course_num(obj):
        return obj.program.program_course.all().count()

    # def get_file(self, obj):
    #     user = self.context.get('user')
    #     if user:
    #         try:
    #             instance = StudentCertificate.objects.get(student_id=user.id, program_id=obj.program.id)
    #             return instance.file
    #         except StudentCertificate.DoesNotExist:
    #             return None
    #     return None

class StudentProgramCertificateSerializer(serializers.ModelSerializer):
    id = serializers.IntegerField(read_only=True)
    student = UserMinSerializer(required=False)
    # program = ProgramMinSerializer(required=False)

    class Meta:
        model = StudentProgram
        fields = [
            'id', 'file', 'student', 'status', 'started_date', 'completed_date', 'received_certificate'
        ]

class ProgramCertificateSerializer(serializers.ModelSerializer):

    id = serializers.IntegerField(read_only=True)
    courses_count = serializers.IntegerField(read_only=True)
    student_program = serializers.SerializerMethodField(read_only=True)

    class Meta:
        model = Program
        fields = [
            'id', 'name', 'icon', 'slug', 'student_program', 'courses_count', 'status', 'field',
        ]

    @staticmethod
    def get_student_program(obj):
        queryset = StudentProgram.objects.filter(program_id=obj.id)
        return StudentProgramCertificateSerializer(instance=queryset, many=True).data


import json
import os
import argparse

def read_json_simple(json_file, unsolvable_only, exclude=[]):
    # print('Reading file ...')
    with open(json_file) as data_file:
        data = json.load(data_file)
    # print('Grouping data ...')
    excluded = 0
    grouped_data = {}
    existing_problems = {}
    for idx, val in data.items():
        if unsolvable_only and val['unsolvable'] == 0:
            continue
        if val['id'][1] not in grouped_data:
            existing_problems[val['id'][1]] = set()
            grouped_data[val['id'][1]] = {}
        if val['id'][1] in grouped_data and val['id'][2] not in grouped_data[val['id'][1]]:
            existing_problems[val['id'][1]].add(val['id'][2])
            grouped_data[val['id'][1]][val['id'][2]] = {}
        if val['id'][0] not in exclude:
            grouped_data[val['id'][1]][val['id'][2]][val['id'][0]] = val
        else:
            excluded += 1

    if excluded == 0 and len(exclude) > 0:
        print('Warning: no data excluded for exclude=', str(exclude))

    # print('Filtering intersection only data ...')

    counter = {}
    max_counter = 0
    for domain, problems in grouped_data.items():
        counter[domain] = {}
        for problem, algos in problems.items():
            counter[domain][problem] = len(algos)
            max_counter = max(len(algos), max_counter)
    to_del = []
    for domain, problems in grouped_data.items():
        for problem, algos in problems.items():
            if counter[domain][problem] < max_counter:
                to_del.append((domain, problem))
    for domain, problem in to_del:
        del grouped_data[domain][problem]
    to_del = []
    for domain, problems in grouped_data.items():
        if len(grouped_data[domain]) == 0:
            to_del.append(domain)
    for domain in to_del:
        del grouped_data[domain]

    # filtering unused instances

    domain_to_del = []
    for domain, problems in grouped_data.items():
        prob_to_del = []
        for problem, algos in problems.items():
            if len(algos) == 0:
                prob_to_del.append(problem)
        for p in prob_to_del:
            del problems[p]
            existing_problems[domain].remove(p)
        if len(problems) == 0:
            domain_to_del.append(domain)
    for d in domain_to_del:
        del grouped_data[d]
        del existing_problems[d]

    return grouped_data, existing_problems

def print_stats(grouped_data, algorithms, attr, max_val):
    # print('Processing data ...')

    # processing data
    val_data = {}
    for domain, problems in grouped_data.items():
        for problem, algos in problems.items():
            for algo, val_list in algos.items():
                if algo in algorithms:
                    if algo not in val_data:
                        val_data[algo] = []
                    if val_list['unsolvable'] == 0:
                        val_data[algo].append(max_val)
                    else:
                        val_data[algo].append(val_list[attr])
    total_data = 0
    for alg in algorithms:
        if total_data == 0:
            total_data = len(val_data[alg])
        else:
            assert len(val_data[alg]) == total_data

    s = ''
    for alg in algorithms:
        s += alg + ' '
    print(s)
    for i in range(total_data):
        s = ''
        for alg in algorithms:
            s += str(val_data[alg][i]) + ' '
        print(s)

def create_dirs_if_necessary(graph_data, OUTDIR):
    for domain, problems in graph_data.items():
        if not os.path.exists(OUTDIR + '/' + domain):
            os.makedirs(OUTDIR + '/' + domain)

def check_attribute_exists(grouped_data, attr):
    attrs = []
    for domain, problems in grouped_data.items():
        for problem, algos in problems.items():
            for algo, runs in algos.items():
                attrs = runs.keys()
            break
        break
    if attr in attrs:
        return True
    else:
        print('Attribute ' + attr + ' doesn\'t exists.')
        print('Available attributes are:')
        attrs = sorted(attrs)
        for a in attrs:
            print(a)
        return False

def check_algorithm_exists(grouped_data, alg):
    attrs = []
    for domain, problems in grouped_data.items():
        for problem, algos in problems.items():
            algs = algos.keys()
        break
    if alg in algs:
        return True
    else:
        print('Algorithm ' + alg + ' doesn\'t exists.')
        print('Algorithm attributes are:')
        algs = sorted(algs)
        for a in algs:
            print(a)
        return False

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("json_file", help=".json file containing the lab data")
    parser.add_argument("--attribute", "-a", help="attribute to use", default="max_abstraction_states")
    parser.add_argument("--unsolvable-only", "-u", help="only count the unsolvable instances", dest='unsolvable_only',
                        action='store_true')
    parser.add_argument("--max", '-m', help="maximum default value if it is fail", default='0')
    parser.add_argument("--algos", help="histogram range (int)", nargs='+')
    parser.set_defaults(unsolvable_only=False)
    args = parser.parse_args()

    data, problems = read_json_simple(args.json_file, args.unsolvable_only)
    if check_attribute_exists(data, args.attribute):
        for alg in args.algos:
            if not check_algorithm_exists(data, alg):
                return
        print_stats(data, args.algos, args.attribute, args.max)

if __name__ == '__main__':
    main()