from __future__ import annotations

# fmt: off
import sys  # isort: skip
from pathlib import Path  # isort: skip
ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(ROOT))
# fmt: on

import gc
import json as json_
import os
import re
import shutil
import traceback
from abc import ABC
from pathlib import Path
from time import sleep
from typing import Any, List, Optional, Tuple
from warnings import warn

import nibabel as nib
import numpy as np
from nibabel import Nifti2Image
from numpy import ndarray
from tqdm.contrib.concurrent import process_map

from rmt.constants import ANAT_REGISTERED_SUFFIX
from rmt.constants import DATA_ROOT as DATA
from rmt.constants import (
    EIGS_SUFFIX,
    EXTRACTED_SUFFIX,
    MASK_SUFFIX,
    MNI_REGISTERED_SUFFIX,
    MOTION_CORRECTED_SUFFIX,
    NII_SUFFIX,
    RMT_DIR,
    SLICETIME_SUFFIX,
    UPDATED,
)
from rmt.preprocess.extract import Loadable, RMTComputatable

if os.environ.get("CC_CLUSTER") == "niagara":
    ROOT = Path("/scratch/j/jlevman/dberger/random-matrix-fmri")
    DATA = Path("/scratch/j/jlevman/dberger/random-matrix-fmri/data")
    UPDATED = Path("/scratch/j/jlevman/dberger/random-matrix-fmri/data/updated")
    RMT_DIR = Path("/scratch/j/jlevman/dberger/random-matrix-fmri/data/updated/rmt")

TEMPLATE = DATA / "tpl-MNI152NLin2009cAsym_res-02_desc-brain_T1w.nii.gz"
if not TEMPLATE.exists():
    raise FileNotFoundError(
        f"No registration template found at {TEMPLATE}. "
        f"Run `python {DATA / 'download_template.py'} to download it."
    )

class FmriScan(Loadable):
    def __init__(self, source: Path) -> None:
        super().__init__(source)
        self.source: Path
        self.t1w_source: Path
        self.sid: str
        self.ses: Optional[str]
        self.run: Optional[str]
        self.local_json: Optional[Path]
        self.global_json: Optional[Path]
        self.slicetimes: Optional[List[str]]
        self.repetition_time: float
        self.slicetime_file: Optional[Path]

        self.source = source
        self.sid, self.ses, self.run, self.data = self.parse_source()
        self.local_json, self.global_json = self.find_json_file()
        self.t1w_source = self.find_t1w_file()
        (
            self.slicetimes,
            self.slicetime_file,
            self.repetition_time,
        ) = self.write_out_slicetimes()

    def brain_extract(self, force: bool = False) -> BrainExtracted:
        """Computes a 4D mask for input fMRI.

        Notes
        -----
        Brain-Extracted Data:

        Non-Extracted:
            Parkinsons
            Learning
            Bilinguality
            Depression
            Osteo

        The ANTS computed mask is truly 4D, e.g. the mask at t=1 will not in
        general be identical to the mask at time t=2. We do unforunately need
        this to get a timeseries at each voxel, and so also must define a
        `min_mask` (or max_mask) for some purposes.

        ants.get_mask seems to be single-core, so it is extremely worth
        parallelizing brain extraction across subjects
        """
        import ants
        from ants import image_read
        from nipype.interfaces.fsl import BET

        outfile = self.mask_path
        if outfile.exists():
            if not force:
                return BrainExtracted(self)
            os.remove(outfile)

        cmd = BET()
        cmd.inputs.in_file = str(self.source.resolve())
        cmd.inputs.out_file = str(self.extracted_path.resolve())
        cmd.inputs.output_type = "NIFTI_GZ"
        cmd.inputs.functional = True
        # cmd.inputs.frac = 0.9  # default with functional is 0.3, leaves too much skull
        if "Vigil" in str(self.source):
            cmd.inputs.frac = 0.3  # default with functional is 0.3, leaves too much skull
        if "Learning" in str(self.source):
            cmd.inputs.frac = 0.7  # default with functional is 0.3, leaves too much skull
        if "Osteo" in str(self.source):
            cmd.inputs.frac = 0.5  # default with functional is 0.3, leaves too much skull
        else:
            cmd.inputs.frac = 0.7  # default with functional is 0.3, leaves too much skull
        cmd.inputs.mask = True

        print(f"Computing mask for {self.source}")
        results = cmd.run()
        maskfile = Path(results.outputs.mask_file).resolve()  # type: ignore
        maskfile.rename(Path(str(maskfile).replace("extracted_", "")))
        print(f"Wrote brain mask to {self.mask_path}")
        print(f"Wrote extracted brain to {self.extracted_path}")

        # dunno, but ANTs causing huge memory usage here
        cmd = None
        results = None
        gc.collect()

        print(f"Cleaning up {self.extracted_path} with ANTS get_mask...")
        img = image_read(str(self.extracted_path))
        if ("Vigil" in str(self.source)) or ("Biling" in str(self.source)):
            # Vigilance data is high spatial resolution, and Bilingual data is very
            # long (800 timepoints). Takes too long to use get_mask with the full
            # image data in both cases, so we use the average image instead
            avg = img.ndimage_to_list()[0].new_image_like(img.mean(axis=-1))
            mask3d = avg.get_mask()
            mask4d = np.stack([mask3d.numpy() for _ in range(img.shape[-1])], axis=-1)
            mask = img.astype("uint8").new_image_like(mask4d)
        else:
            mask = img.get_mask()
        img = img.mask_image(mask)
        ants.image_write(img, str(self.extracted_path))
        return BrainExtracted(self)

    def anat_extract(self, force: bool = False) -> AnatExtracted:
        """Computes a 4D mask for input fMRI.

        Notes
        -----
        Using "robust" option causes a huge amount of phony stderr spam about missing or
        truncated files. Ignore it.
        """
        from nipype.interfaces.base.support import InterfaceResult
        from nipype.interfaces.fsl import BET

        outfile = Path(
            str(self.t1w_source.resolve()).replace(NII_SUFFIX, EXTRACTED_SUFFIX)
        )
        maskfile = Path(str(self.t1w_source.resolve()).replace(NII_SUFFIX, MASK_SUFFIX))
        if outfile.exists():
            if not force:
                return AnatExtracted(self, mask=maskfile, extracted=outfile)
            if "ses-2" in str(outfile):
                # special case, no need to recompute even when forcing
                return AnatExtracted(self, mask=maskfile, extracted=outfile)
            os.remove(outfile)

        cmd = BET()
        cmd.inputs.in_file = str(self.t1w_source.resolve())
        cmd.inputs.out_file = str(outfile)
        cmd.inputs.output_type = "NIFTI_GZ"
        if "Depress" in str(self.t1w_source):
            # this dataset is very strange, needs lower frac to maintain more brain...
            # some subjects still lose a bit in some slices with the lower value below,
            # but with frac = 0.05 only really 2 subjects are very mangled in some slices
            cmd.inputs.frac = 0.05
        elif "Older" in str(self.t1w_source):
            # This dataset has a lot of neck left behind for some reason, so we use the
            # robust option. Experimentation also hows that the `frac` below works best
            cmd.inputs.robust = True
            cmd.inputs.frac = 0.7
        elif "Vigil" in str(self.t1w_source):
            # This also needs the robust treatment. Huge amounts of neck.
            # Also there are only 22 subjects for this dataset
            cmd.inputs.robust = True
            cmd.inputs.frac = 0.3
        else:
            cmd.inputs.frac = 0.3
        cmd.inputs.mask = True

        print(f"Computing mask for {self.t1w_source}")
        results: InterfaceResult = cmd.run()
        bet_maskfile = Path(results.outputs.mask_file).resolve()  # type: ignore
        sleep(1)
        if bet_maskfile.exists():
            shutil.move(bet_maskfile, maskfile)
        else:
            from pprint import pformat

            raise FileNotFoundError(
                f"\n===========================\n"
                f"\tCannot find maskfile: {bet_maskfile}. Details:\n"
                f"\tresults.inputs:\n{pformat(results.inputs, indent=2, depth=5)}\n"
                f"\tresults.outputs:\n{pformat(results.outputs, indent=2, depth=5)}\n"
                f"\toutfile: {outfile}\n"
                f"\tmaskfile: {maskfile}\n"
                f"===========================\n"
            )
        print(f"Wrote brain mask to {maskfile}")
        print(f"Wrote extracted brain to {outfile}")
        return AnatExtracted(self, mask=maskfile, extracted=outfile)

    def reorient_4d_to_RAS(self, img: Any) -> Any:
        import ants
        from ants import ANTsImage

        assert isinstance(img, ANTsImage)

        arrs = [ants.reorient_image2(im).numpy() for im in ants.ndimage_to_list(img)]
        data = np.stack(arrs, axis=-1)
        reoriented = img.new_image_like(data)
        return reoriented

    @property
    def mask_path(self) -> Path:
        parent = self.source.parent
        stem = str(self.source.resolve().name).replace(NII_SUFFIX, "")
        outname = f"{stem}_mask.nii.gz"
        return Path(parent / outname)

    @property
    def extracted_path(self) -> Path:
        parent = self.source.parent
        outname = str(self.source.resolve().name).replace(NII_SUFFIX, EXTRACTED_SUFFIX)
        return Path(parent / outname)

    def parse_source(self) -> Tuple[str, Optional[str], Optional[str], Optional[str]]:
        sid_ = re.search(r"sub-(?:[a-zA-Z]*)?([0-9]+)_.*", self.source.stem)
        ses = re.search(r"ses-([0-9]+)_.*", self.source.stem)
        run_ = re.search(r"run-([0-9]+)_.*", self.source.stem)
        data_ = re.search(r"data/updated/(.*)/ds00.*", str(self.source))

        if sid_ is None:
            raise RuntimeError(f"Didn't find an SID in filename {self.source}")
        sid = sid_[1]
        session = ses[1] if ses is not None else None
        run = run_[1] if run_ is not None else None
        data = data_[1] if data_ is not None else None
        return sid, session, run, data

    def find_t1w_file(self) -> Path:
        anat_dir = self.source.parent.parent / "anat"
        if "ses-2" in str(anat_dir):
            anat_dir = Path(str(anat_dir).replace("ses-2", "ses-1"))
        # if anat_dir.parent.name == "ses-2":
        #     anat_dir = anat_dir.parent.parent / "ses-1/anat"
        anat_img = sorted(anat_dir.glob("*T1w.nii.gz"))
        if len(anat_img) == 0:
            raise RuntimeError(f"Missing T1w file at {anat_dir}")
        if len(anat_img) > 1:
            raise RuntimeError(f"Too many T1w files at {anat_dir}")

        return anat_img[0]

    def find_json_file(self) -> Tuple[Optional[Path], Optional[Path]]:
        # recurse up to BIDS root for remaining json files
        root = self.source.parent
        while "download" not in root.name:
            root = root.parent

        local_json: Optional[Path] = Path(str(self.source).replace(NII_SUFFIX, ".json"))
        if not local_json.exists():  # type: ignore
            local_json = None

        # just hardcode the remaining cases where data is not per scan:
        source = self.source.name

        # this data has local files but without slice timing in them
        if "Park" in str(root):
            global_json: Optional[Path] = root / "task-ANT_bold.json"
            return local_json, global_json

        # Rest_w_VigilanceAttention
        if "prefrontal_bold" in source:
            global_json = root / "task-rest_acq-prefrontal_bold.json"
        elif "rest_acq-fullbrain" in source:
            global_json = root / "task-rest_acq-fullbrain_bold.json"
        # Park_v_Control
        elif ("RC" in source) and ("task-ANT" in source):
            global_json = root / "task-ANT_bold.json"
        # Rest_w_Depression_v_Control
        elif "Depression" in root.parent.name:
            global_json = root / "task-rest_bold.json"
        # Rest_w_Healthy_v_OsteoPain
        elif "Osteo" in root.parent.name:
            global_json = root / "task-rest_bold.json"
        elif "Learning" in str(root):
            global_json = None
        elif "Biling" in str(root):
            global_json = None
        elif "Older" in str(root):
            global_json = None
        else:
            raise RuntimeError(f"Can't find .json info from source: {self.source}")
        return local_json, global_json

    def write_out_slicetimes(self) -> Tuple[Optional[List[str]], Optional[Path], float]:
        timings: Optional[List[str]] = None
        TR: Optional[float] = None
        for jsonfile in [self.local_json, self.global_json]:
            if jsonfile is None:
                continue

            with open(jsonfile, "r") as handle:
                info = json_.load(handle)
            found_timings: Optional[List[str]] = info.get("SliceTiming", None)
            if (found_timings is not None) and (timings is None):
                timings = found_timings

            found_TR = info.get("RepetitionTime", None)
            if (found_TR is not None) and (TR is None):
                TR = float(found_TR)

        if TR is None:
            raise KeyError(
                f"Could not find RepetitionTime field in any .json for {self.source}"
            )

        if timings is not None:
            outfile: Optional[Path] = Path(
                str(self.source).replace(NII_SUFFIX, ".slicetimes.txt")
            )
            if not outfile.exists():  # type: ignore
                with open(outfile, "w") as handle:  # type: ignore
                    handle.writelines(list(map(lambda t: f"{t}\n", timings)))
                print(f"Wrote slice timings to {outfile}")
        else:
            outfile = None

        return timings, outfile, TR

class AnatExtracted(Loadable):
    def __init__(self, raw: FmriScan, mask: Path, extracted: Path) -> None:
        super().__init__(extracted)
        self.raw = raw
        self.mask = mask
        self.source: Path = extracted

class BrainExtracted(RMTComputatable):
    def __init__(self, raw: FmriScan) -> None:
        super().__init__(raw.extracted_path)
        self.raw = raw
        self.mask = self.raw.mask_path
        self.source: Path = raw.extracted_path

    def slicetime_correct(
        self, slice_direction: int = 3, force: bool = False
    ) -> SliceTimeCorrected:
        """
        Notes
        -----
        Only Rest_w_VigilanceAttention data has SliceEncodingDirection = "k"
        (i.e. k+, slices along third spatial dimension, first entry of file
        corresponds to smallest index along thid spatial dim)

        For BIDS slice timing documentation, see:

        https://poc.vl-e.nl/distribution/manual/fsl-3.2/slicetimer/index.html

        reproduced below:

        # INTRODUCTION

            slicetimer is a pre-processing tool designed to correct for sampling
            offsets inherent in slice-wise EPI acquisition sequences.

            Each voxel's timecourse is processed independently and intensities
            are shifted in time so that they reflect the interpolated value of
            the signal at a common reference timepoint for all voxels, providing
            an instantaneous `snapshot' of the data, rather than a staggered
            sample throughout each volume. Sinc interpolation with a Hanning
            windowing kernel is applied to each timecourse to calculate the
            interpolated values.

            It is necessary to know in what order the slices were acquired and
            set the appropriate option. The default correction is appropriate if
            slices were acquired from the bottom of the brain.

            If slices were acquired from the top of the brain to the bottom
            select the --down option.

            If the slices were acquired with interleaved order (0, 2, 4 ... 1,
            3, 5 ...) then choose the --odd option.

            If slices were not acquired in regular order you will need to use a
            slice order file or a slice timings file. If a slice order file is
            to be used, create a text file with a single number on each line,
            where the first line states which slice was acquired first, the
            second line states which slice was acquired second, etc. The first
            slice is numbered 1 not 0.

            If a slice timings file is to be used, put one value (ie for each
            slice) on each line of a text file. The units are in TRs, with 0.5
            corresponding to no shift. Therefore a sensible range of values will
            be between 0 and 1.

        #  USAGE AND OPTIONS

            ```
            slicetimer -i  [-o ] [options]

            Compulsory arguments (You MUST set one or more of):

                    -i,--in filename of input timeseries

            Optional arguments (You may optionally specify one or more of):

                    -o,--out        filename of output timeseries
                    -h,--help       display this message
                    -v,--verbose    switch on diagnostic messages
                    --down          reverse slice indexing
                    -r,--repeat     Specify TR of data - default is 3s
                    -d,--direction  direction of slice acquisition (x=1,y=2,z=3)
                                    default is z
                    --odd           use interleaved acquisition
                    --tcustom       filename of single-column custom interleave timing
                    --ocustom       filename of single-column custom interleave order
                                    (first slice is referred to as 1 not 0)
            ```
        """
        from nipype.interfaces.fsl import SliceTimer

        infile = self.source
        outfile = Path(str(infile).replace(NII_SUFFIX, SLICETIME_SUFFIX))
        if outfile.exists():
            if not force:
                return SliceTimeCorrected(self, outfile)
            os.remove(outfile)

        if (self.raw.slicetime_file is None) or (not self.raw.slicetime_file.exists()):
            # just copy over extracted brain
            shutil.copy(infile, outfile)
            warn(
                f"No slicetimes for {self.source}. Copying brain-extracted file instead."
            )
            return SliceTimeCorrected(self, outfile)

        cmd = SliceTimer()
        cmd.inputs.in_file = str(infile)
        cmd.inputs.custom_timings = str(self.raw.slicetime_file)
        cmd.inputs.time_repetition = self.raw.repetition_time
        cmd.inputs.out_file = str(outfile)
        cmd.inputs.output_type = "NIFTI_GZ"
        cmd.inputs.slice_direction = slice_direction
        print(f"Slice time correcting {infile}...")
        cmd.run()
        return SliceTimeCorrected(self, outfile)

class SliceTimeCorrected(RMTComputatable):
    def __init__(self, extracted: BrainExtracted, corrected: Path) -> None:
        super().__init__(corrected)
        self.raw: FmriScan = extracted.raw
        self.extracted: BrainExtracted = extracted
        self.source: Path = corrected

    def eigenvalues(self) -> ndarray:
        """Extract masked correlation matrix from self.source and compute eigenvalues"""
        raise NotImplementedError()

    def motion_corrected(self, force: bool = False) -> MotionCorrected:
        import ants
        from ants import motion_correction

        outfile = Path(str(self.source).replace(NII_SUFFIX, MOTION_CORRECTED_SUFFIX))
        if outfile.exists():
            if not force:
                return MotionCorrected(self, outfile)
            os.remove(outfile)

        img = ants.image_read(str(self.source))
        # img8 = img.ndimage_to_list()[8]  # don't use first volume...
        # mask3d = img8.get_mask()
        print(f"Finding transform for {self.source}")
        results = motion_correction(
            img,
            fixed=None,  # uses mean image in this case
            type_of_transform="BOLDRigid",
            # mask=mask3d,
        )
        corrected = results["motion_corrected"]
        print(f"Performing motion correction for {self.source}")
        ants.image_write(corrected, str(outfile))
        print(f"Saved motion-corrected image to {outfile}")
        return MotionCorrected(self, motion_corrected=outfile)

class MotionCorrected(RMTComputatable):
    def __init__(self, corrected: SliceTimeCorrected, motion_corrected: Path) -> None:
        self.raw = corrected.raw
        self.extracted = corrected.extracted
        self.slicetime_corrected: SliceTimeCorrected = corrected
        self.source: Path = motion_corrected
        super().__init__(self.source)

    @staticmethod
    def reorient_template_to_img(template: Any, img: Any) -> Any:
        """The actual ANTs functions are completely broken for some reason, so we
        do it manually...

        Notes
        -----
                               data  x_n  y_n   z_n x_mm  y_mm  z_mm    t       TR orient
                     Park_v_Control   80   80   43  3.00  3.00  3.00  149     2.40    RPI
                     Park_v_Control   80   80   43  3.00  3.00  3.00  300     2.40    RPI
                     Park_v_Control   96  114   96  2.00  2.00  2.00  149     2.40    LPI
              Rest_v_LearningRecall   64   64   36  3.00  3.00  3.00  195     2.00    RPI
               Rest_w_Bilinguiality  100  100   72  1.80  1.80  1.80  823     0.88    RPI
               Rest_w_Bilinguiality  100   96   72  1.80  1.80  1.80  823     0.88    RPI
               Rest_w_Bilinguiality  100  100   72  1.80  1.80  1.80  823     0.93    RIA
        Rest_w_Depression_v_Control  112  112   25  1.96  1.96  5.00  100     2.50    RPI
         Rest_w_Healthy_v_OsteoPain   64   64   36  3.44  3.44  3.00  244     2.50    RPI
         Rest_w_Healthy_v_OsteoPain   64   64   36  3.44  3.44  3.00  292     2.50    RPI
         Rest_w_Healthy_v_OsteoPain   64   64   36  3.44  3.44  3.00  300     2.50    RPI
             Rest_w_Older_v_Younger   74   74   32  2.97  2.97  4.00  300     2.00    RPI
          Rest_w_VigilanceAttention   64   64   35  3.00  3.00  3.00  300  3000.00    RPI
          Rest_w_VigilanceAttention  128  128   70  1.50  1.50  1.50  300  3000.00    RPI
          Rest_w_VigilanceAttention  200   60   40  0.75  0.75  0.75  150  4000.00    RPI

        >>> df.value_counts()
        data                      x_n  y_n  z_n  x_mm  y_mm  z_mm  t    TR       orient
        Park_v_Control            80   80   43   3.00  3.00  3.00  149  2.40     RPI   552
                                  96   114  96   2.00  2.00  2.00  149  2.40     LPI   552
        Rest_v_LearningRecall     64   64   36   3.00  3.00  3.00  195  2.00     RPI   432
        Rest_w_Bilinguiality      100  100  72   1.80  1.80  1.80  823  0.88     RPI    90
        Rest_w_VigilanceAttenti   128  128  70   1.50  1.50  1.50  300  3000.00  RPI    84
        Rest_w_Healthy_v_OsteoP   64   64   36   3.44  3.44  3.00  300  2.50     RPI    74
        Rest_w_Depression_v_Conl  112  112  25   1.96  1.96  5.00  100  2.50     RPI    72
        Rest_w_Older_v_Younger    74   74   32   2.97  2.97  4.00  300  2.00     RPI    62
        Rest_w_VigilanceAttenti   200  60   40   0.75  0.75  0.75  150  4000.00  RPI    44
                                  64   64   35   3.00  3.00  3.00  300  3000.00  RPI     4
        Park_v_Control            80   80   43   3.00  3.00  3.00  300  2.40     RPI     1
        Rest_w_Bilinguiality      100  96   72   1.80  1.80  1.80  823  0.88     RPI     1
                                       100  72   1.80  1.80  1.80  823  0.93     RIA     1
        Rest_w_Healthy_v_OsteoP   64   64   36   3.44  3.44  3.00  244  2.50     RPI     1
                                                                   292  2.50     RPI     1

        """
        from ants import ANTsImage

        assert isinstance(template, ANTsImage)
        assert isinstance(img, ANTsImage)

        temp_orient = template.get_orientation()
        img_orient = img.get_orientation()
        oriented: ANTsImage = template.clone()
        if (temp_orient == "LPI") and (img_orient == "RPI"):
            return oriented.reflect_image(axis=1).reflect_image(axis=2)
        else:
            raise ValueError(
                "Impossible! All scans to register should be RPI orientation"
            )

        # for i, (axis_temp, axis_img) in enumerate(zip(temp_orient, img_orient)):
        #     if axis_temp == axis_img:
        #         continue
        #     oriented = oriented.reflect_image(axis=i)
        # # reoriented = template.new_image_like(oriented)
        # reoriented = oriented
        # return reoriented

    def t1w_register(self, force: bool = False) -> T1wRegistered:
        """
        Notes
        -----
        fMRI is too low-q to register straight to MNI template. we need to see if
        it is better to go through anat in one way or another. e.g. the classic

            fMRI: reg to -> ant: reg to -> MNI

        or:

            reg fMRI to anat: save reverse transform (anat_to_fmri)
            reg anat to MNI: save reverse transform (mni_to_anat)

        # State of Extraction

        Not extracted:
            Park_v_Control
            Bilingual
            Vigilance
            Osteo
            Learning
            Older_v_Younger

        Extracted:
            Depression

        """
        import ants
        from ants import ANTsImage

        outfile = Path(str(self.source).replace(NII_SUFFIX, ANAT_REGISTERED_SUFFIX))
        if outfile.exists():
            if not force:
                return T1wRegistered(self, registered=outfile)
            os.remove(outfile)

        img = ants.image_read(str(self.source))
        imgs: List[ANTsImage] = ants.ndimage_to_list(img)
        avg = imgs[0].new_image_like(img.mean(axis=-1))
        avg_mask = avg.get_mask()
        anat_cls = self.raw.anat_extract(force=False)
        anat = ants.image_read(str(anat_cls.source))

        # having the moving mask (`mask=avg_mask` below) is crucial to preventing
        # severe distortions in the final registered image due to missing regions
        print(f"Registering {self.source} average image to {str(anat_cls.source)}")
        results = ants.registration(
            fixed=anat,
            moving=avg,
            mask=avg_mask,
            type_of_transform="SyNBold",
        )
        transforms = results["fwdtransforms"]
        print("Applying transform from average image to full 4D data")
        img = ants.apply_transforms(
            fixed=anat,
            moving=img,
            transformlist=transforms,
            imagetype=3,
        )
        ants.image_write(img, str(outfile))
        print(f"Saved T1w-registered image to {outfile}")

        # Save quality check image while we are here
        plot_out = Path(str(outfile).replace(NII_SUFFIX, "_T1w-reg-plot.png"))
        avg_reg = img.ndimage_to_list()[0].new_image_like(img.mean(axis=-1))
        avg_reg.plot(filename=str(plot_out), reorient=False)
        return T1wRegistered(self, registered=outfile)

    def mni_register(self, force: bool = False) -> MNI152Registered:
        """Register fMRI directly (!!) to MNI 2x2x2 mm template"""
        import ants
        from ants import ANTsImage

        outfile = Path(str(self.source).replace(NII_SUFFIX, MNI_REGISTERED_SUFFIX))
        if outfile.exists():
            if not force:
                return MNI152Registered(self, registered=outfile)
            os.remove(outfile)

        img = ants.image_read(str(self.source))
        imgs: List[ANTsImage] = ants.ndimage_to_list(img)
        avg = imgs[0].new_image_like(img.mean(axis=-1))
        avg_mask = avg.get_mask()
        template = ants.image_read(str(TEMPLATE))

        # having the moving mask (`mask=avg_mask` below) is crucial to preventing
        # severe distortions in the final registered image due to missing regions
        print(f"Registering {self.source} average image to {TEMPLATE}")
        results = ants.registration(
            fixed=template,
            moving=avg,
            mask=avg_mask,
            type_of_transform="SyNBold",
        )
        transforms = results["fwdtransforms"]
        print("Applying transform from average image to full 4D data")
        img = ants.apply_transforms(
            fixed=template,
            moving=img,
            transformlist=transforms,
            imagetype=3,
        )
        ants.image_write(img, str(outfile))
        print(f"Saved MNI-registered image to {outfile}")

        # Save quality check image while we are here
        plot_out = Path(str(outfile).replace(NII_SUFFIX, "_reg-plot.png"))
        avg_reg = img.ndimage_to_list()[0].new_image_like(img.mean(axis=-1))
        avg_reg.plot(filename=str(plot_out), reorient=False)
        return MNI152Registered(self, registered=outfile)

class T1wRegistered(Loadable):
    def __init__(self, motion_corrected: MotionCorrected, registered: Path) -> None:
        self.raw = motion_corrected.raw
        self.extracted = motion_corrected.extracted
        self.slicetime_corrected: SliceTimeCorrected = (
            motion_corrected.slicetime_corrected
        )
        self.motion_corrected: MotionCorrected = motion_corrected
        self.source: Path = registered
        super().__init__(self.source)

class MNI152Registered(RMTComputatable):
    def __init__(self, motion_corrected: MotionCorrected, registered: Path) -> None:
        self.raw = motion_corrected.raw
        self.extracted = motion_corrected.extracted
        self.slicetime_corrected: SliceTimeCorrected = (
            motion_corrected.slicetime_corrected
        )
        self.motion_corrected: MotionCorrected = motion_corrected
        self.source: Path = registered
        super().__init__(self.source)

def get_fmri_paths(filt: Optional[str] = None) -> List[Path]:
    UPDATED = DATA / "updated"
    parents = sorted(filter(lambda p: p.is_dir(), UPDATED.glob("*")))
    paths = []
    for parent in parents:
        paths.extend(sorted(parent.rglob("*bold.nii.gz")))
    paths = sorted(filter(lambda p: "derivative" not in str(p), paths))
    paths = sorted(filter(lambda p: "prefrontal" not in str(p), paths))
    if filt is not None:
        paths = sorted(filter(lambda p: filt in str(p), paths))  # type: ignore
    return paths

def brain_extract_parallel(path: Path) -> None:
    try:
        fmri = FmriScan(path)
        fmri.brain_extract(force=True)
    except Exception:
        traceback.print_exc()

def anat_extract_parallel(path: Path) -> None:
    try:
        fmri = FmriScan(path)
        fmri.anat_extract(force=True)
    except Exception:
        traceback.print_exc()

def make_slicetime_file(path: Path) -> None:
    try:
        FmriScan(path)
    except Exception:
        traceback.print_exc()

def inspect_extractions(path: Path) -> None:
    from ants import image_read

    try:
        fmri = FmriScan(path)
        stripped = fmri.brain_extract(force=False)
        orig_file = str(path).replace(NII_SUFFIX, "_plot.png")
        extr_file = str(stripped.source).replace(NII_SUFFIX, "_plot.png")
        if Path(extr_file).exists():
            return
        orig = image_read(str(path))
        extracted = image_read(str(stripped.source))
        orig.ndimage_to_list()[5].plot(filename=orig_file)
        extracted.ndimage_to_list()[5].plot(filename=extr_file)
    except Exception:
        traceback.print_exc()

def reinspect_extractions(path: Path) -> None:
    from ants import image_read

    try:
        fmri = FmriScan(path)
        stripped = fmri.brain_extract(force=False)
        orig_file = str(path).replace(NII_SUFFIX, "_plot.png")
        extr_file = str(stripped.source).replace(NII_SUFFIX, "_plot.png")
        orig = image_read(str(path))
        extracted = image_read(str(stripped.source))
        orig.ndimage_to_list()[5].plot(filename=orig_file)
        extracted.ndimage_to_list()[5].plot(filename=extr_file)
    except Exception:
        traceback.print_exc()

def inspect_anat_extractions(path: Path) -> None:
    from ants import image_read

    try:
        fmri = FmriScan(path)
        stripped = fmri.anat_extract(force=False)
        orig_file = str(fmri.t1w_source).replace(NII_SUFFIX, "_plot.png")
        extr_file = str(stripped.source).replace(NII_SUFFIX, "_plot.png")
        if Path(extr_file).exists():
            return
        orig = image_read(str(fmri.t1w_source))
        extracted = image_read(str(stripped.source))
        orig.plot(filename=orig_file)
        extracted.plot(filename=extr_file)
    except Exception:
        traceback.print_exc()

def reinspect_anat_extractions(path: Path) -> None:
    from ants import image_read

    try:
        fmri = FmriScan(path)
        stripped = fmri.anat_extract(force=False)
        orig_file = str(fmri.t1w_source).replace(NII_SUFFIX, "_plot.png")
        extr_file = str(stripped.source).replace(NII_SUFFIX, "_plot.png")
        orig = image_read(str(fmri.t1w_source))
        extracted = image_read(str(stripped.source))
        orig.plot(filename=orig_file)
        extracted.plot(filename=extr_file)
    except Exception:
        traceback.print_exc()

def slicetime_correct_parallel(path: Path) -> None:
    try:
        fmri = FmriScan(path)
        extracted = fmri.brain_extract(force=False)
        extracted.slicetime_correct(force=False)
    except Exception:
        traceback.print_exc()

def motion_correct_parallel(path: Path) -> None:
    try:
        fmri = FmriScan(path)
        extracted = fmri.brain_extract(force=False)
        corrected = extracted.slicetime_correct(force=False)
        corrected.motion_corrected(force=False)
    except Exception:
        traceback.print_exc()

def mni_register_parallel(path: Path) -> None:
    try:
        fmri = FmriScan(path)
        extracted = fmri.brain_extract(force=False)
        slice_corrected = extracted.slicetime_correct(force=False)
        motion_corrected = slice_corrected.motion_corrected(force=False)
        motion_corrected.mni_register(force=False)
    except Exception:
        traceback.print_exc()

if __name__ == "__main__":
    # on Niagara need module load gcc/8.3.0 openblas/0.3.7 fsl/6.0.4
    # paths = get_fmri_paths(filt="Park")
    paths = get_fmri_paths()
    # process_map(make_slicetime_file, paths, chunksize=1)
    # NOTE: for "Vigilance" data, can only have up to 10 workers
    # NOTE: for "Bilingual" data, can only have up to 7? workers
    # process_map(brain_extract_parallel, paths, max_workers=40)
    # process_map(inspect_extractions, paths, chunksize=1, max_workers=40)
    # process_map(anat_extract_parallel, paths, chunksize=1, max_workers=40)
    # process_map(inspect_extractions, paths, chunksize=1, max_workers=40)
    # process_map(inspect_anat_extractions, paths, chunksize=1, max_workers=40)
    # process_map(reinspect_anat_extractions, paths, chunksize=1, max_workers=40)
    # process_map(slicetime_correct_parallel, paths, chunksize=1, max_workers=40)
    # process_map(motion_correct_parallel, paths, max_workers=8)
    # process_map(reinspect_extractions, paths, chunksize=1, max_workers=40)
    process_map(mni_register_parallel, paths, max_workers=12)


import logging
import os.path
from dataclasses import dataclass
from typing import Iterator, List, Optional

from common import directories, file_utils
from common.commands.base import ArxivBatchCommand
from common.types import ArxivId

from ..types import AnnotationToken, Token

@dataclass(frozen=True)
class PaperTokens:
    arxiv_id: ArxivId
    tokens: List[Token]

class CreateAnnotationFiles(ArxivBatchCommand[PaperTokens, None]):
    @staticmethod
    def get_name() -> str:
        return "create-annotation-files"

    @staticmethod
    def get_description() -> str:
        return (
            "Generate files for annotation with an annotation tool (e.g., 'brat') from tokens "
            + "extracted from TeX files. For each paper, three files are generated: (1) A cleaned "
            + "up text meant to be annotated; (2) an updated list of tokens containing position "
            + "information about where each one appears in the cleaned text; (3) an annotation "
            + "file for 'brat' containing a list of automatically generated initial annotations."
        )

    def get_arxiv_ids_dirkey(self) -> str:
        return "sentence-tokens"

    def load(self) -> Iterator[PaperTokens]:
        for arxiv_id in self.arxiv_ids:

            output_dir = directories.arxiv_subdir("annotation-files", arxiv_id)
            file_utils.clean_directory(output_dir)

            # Load tokens.
            tokens_path = os.path.join(
                directories.arxiv_subdir("sentence-tokens", arxiv_id), "tokens.csv",
            )
            try:
                tokens = list(file_utils.load_from_csv(tokens_path, Token))
            except FileNotFoundError:
                logging.warning(  # pylint: disable=logging-not-lazy
                    "No tokens data found for arXiv paper %s. No annotation files will be "
                    + "generated for this paper.",
                    arxiv_id,
                )
                continue

            yield PaperTokens(arxiv_id, tokens)

    def process(self, _: PaperTokens) -> Iterator[None]:
        yield None

    def _create_annotation_token(
        self,
        token: Token,
        start: Optional[int],
        end: Optional[int],
        start_in_sentence: Optional[int],
        end_in_sentence: Optional[int],
    ) -> AnnotationToken:
        return AnnotationToken(
            text=token.text,
            text_for_annotation=token.text_for_annotation,
            tex=token.tex,
            tex_start=token.tex_start,
            tex_end=token.tex_end,
            sentence_id=token.sentence_id,
            sentence_transformations=token.sentence_transformations,
            start_in_text=start,
            end_in_text=end,
            start_in_sentence=start_in_sentence,
            end_in_sentence=end_in_sentence,
        )

    def save(self, item: PaperTokens, _: None) -> None:

        output_dir = directories.arxiv_subdir("annotation-files", item.arxiv_id)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # Write tokens simultaneously to a file listing each token, and a file including
        # the transformed, cleaned text ready for human annotation.
        token_list_path = os.path.join(output_dir, "tokens.csv")
        sentence_list_path = os.path.join(output_dir, "sentences_for_annotation.txt")
        annotation_file_path = os.path.join(output_dir, "initial_annotations.ann")

        offset_in_file = 0
        offset_in_sentence = 0
        is_symbol = False
        brat_annotation_index = 1

        with open(
            sentence_list_path, mode="w", encoding="utf-8"
        ) as sentence_file, open(
            annotation_file_path, mode="w", encoding="utf-8"
        ) as ann_file:
            last_sentence_id = None

            for token in item.tokens:

                # Keep track of whether an upcoming token is a symbol. If it is, a tentative
                # annotation can be saved to a brat annotation file.
                if not is_symbol and token.text == "SYMBOL_START":
                    is_symbol = True
                elif is_symbol and token.text == "SYMBOL_END":
                    is_symbol = False

                # Skip tokens that lack a clean readable representation.
                if token.text_for_annotation is None:
                    annotation_token = self._create_annotation_token(
                        token,
                        start=None,
                        end=None,
                        start_in_sentence=None,
                        end_in_sentence=None,
                    )
                    file_utils.append_to_csv(token_list_path, annotation_token)
                    continue

                # Break lines between sentences.
                if (
                    token.sentence_id != last_sentence_id
                    and last_sentence_id is not None
                ):
                    sentence_file.write("\n")
                    offset_in_file += 1
                    offset_in_sentence = 0

                # Insert spaces between tokens.
                if token.sentence_id == last_sentence_id:
                    sentence_file.write(" ")
                    offset_in_file += 1
                    offset_in_sentence += 1

                start = offset_in_file
                start_in_sentence = offset_in_sentence

                # Write token to the text meant for annotation.
                sentence_file.write(token.text_for_annotation)
                offset_in_file += len(token.text_for_annotation)
                offset_in_sentence += len(token.text_for_annotation)

                end = offset_in_file
                end_in_sentence = offset_in_sentence
                last_sentence_id = token.sentence_id

                # Save a record of this token and its positions in the TeX and annotation file.
                annotation_token = self._create_annotation_token(
                    token,
                    start=start,
                    end=end,
                    start_in_sentence=start_in_sentence,
                    end_in_sentence=end_in_sentence,
                )
                file_utils.append_to_csv(token_list_path, annotation_token)

                # Save an annotation for each symbol. Annotations are stored to a '.ann' file in
                # brat standoff format. See more details about the format here:
                # https://brat.nlplab.org/standoff.html
                if is_symbol:
                    ann_file.write(
                        f"T{brat_annotation_index}\t"
                        + f"Symbol {start} {end}\t"
                        + f"{token.text_for_annotation}\n"
                    )
                    brat_annotation_index += 1


import open3d as o3d
import numpy as np

default_colors = np.array([ [0,0,0],
                            [255,255,255],
                            [0,101,189],
                            [0,82,147],
                            [100,160,200],
                            [152,198,234],
                            [153,153,153],
                            [227,114,34],
                            [162,173,0],
                            [218,215,203],
                            [0,90,90],
                            [90,90,0]    ])

def available_methodes():
    return ['bbd',
            'v2v',
            'IoU_v',
            'IoU_p',
            'pd',
            'dd',
            'od_e',
            'od_R'
            ]

def positionDifference(pos1,pos2):
    return np.linalg.norm(pos1-pos2)

def dimensionDifference(dim1,dim2):
    return abs(dim1.prod()-dim2.prod())  # one could also use l2 norm instead of volume

def circDifference_euler(rot1,rot2,circ=360):
    from scipy.spatial.transform import Rotation as R
    ang1 = R.from_matrix(rot1).as_euler('xyz')
    ang2 = R.from_matrix(rot2).as_euler('xyz')
    return sum(abs(np.mod(ang1-ang2 + circ/2, circ) - circ/2)) # one could also use l2 norm instead of l1

def rotDifference(R1,R2):
    return np.arccos((np.trace(R1@R2.T)-1)/2) # transpose of a rot-matrix is its inverse

def computeDist(a,b):
    n=np.zeros([len(a),len(b)])
    for i,p in enumerate(a):
        diff = b-p
        diff = np.linalg.norm(diff,axis=1)
        n[i,:] = diff
    return n

def edge2edgeDist(lines1,lines2, check=True, eps = 1e-8):
    m1 = (lines1[:,1]-lines1[:,0]).T
    m2 = (lines2[:,1]-lines2[:,0]).T

    #m1 = m1# / np.linalg.norm(m1,axis=0)
    #m2 = m2# / np.linalg.norm(m2,axis=0)
    e = (lines1[:,0] - lines2[:,0]).T

    det = np.diag(- (m1.T@m1) * (m2.T@m2) + (m1.T@m2)**2) # = valid
    with np.errstate(divide='ignore',invalid='ignore'):
        t2 = np.diag((-m1.T@m1) * (m2.T@e) + (m1.T@e) * (m1.T@m2)) / det
        t1 = np.diag((m2.T@m2) * (m1.T@e) - (m2.T@e) * (m1.T@m2)) / det

    P1 = lines1[:,0].T + m1 * t1
    P2 = lines2[:,0].T + m2 * t2

    valid = (det!=0)
    valid = valid & (0-eps<=t1) & (t1<=1+eps) & (0-eps<=t2) & (t2<=1+eps) if check else valid
    P1 = P1[:,valid]
    P2 = P2[:,valid]

    dist = np.linalg.norm(P1-P2,axis=0)

    return list(zip(dist,P1.T,P2.T))

def line2pointDist(line,points,check=True,eps = 1e-8):
    m = line[1] - line[0]
    t = (points - line[0])@m / (m@m)
    P = line[0] + (np.expand_dims(m, axis=1) * t).T

    if check:
        valid = (0-eps<=t) & (t<=1+eps)
        P = P[valid]
        points = points[valid]

    dist = np.linalg.norm(P-points,axis=1)

    return list(zip(dist,P,points))

class plane():
    def __init__(self,p0,n):
        self.p0 = p0
        self.n = n/np.linalg.norm(n)

    def __init__(self,p0,p1,p2):
        self.p0 = p0
        self.p1 = p1
        self.p2 = p2
        self.v = np.stack([p1-p0,p2-p0]).T
        n = np.cross(self.v[:,0],self.v[:,1])
        self.n = n/np.linalg.norm(n)

    def intersect_lines(self, lines, eps=1e-8):
        """Finds intersections with lines given by tuples"""
        m = lines[:,1]-lines[:,0] # gradient
        dot = self.n @ m.T # angles between plane and lines

        valid = abs(dot) > eps # only non parallel lines have valid solutions

        t = (self.n @ (self.p0 - lines[valid,0]).T) / dot[valid] # line parameter
        intersections = lines[valid,0] + (m[valid]*t[..., np.newaxis])

        return intersections, t

    def project_points(self, points, check=True, eps = 1e-8):
        v = (points-self.p0).T
        dist = self.n @ v
        prj_points = points - self.n * dist[..., np.newaxis]

        if check:
            t = np.linalg.inv(self.v.T @ self.v) @ self.v.T @ v
            valid = (0-eps<=t[0]) & (t[0]<=1+eps) & (0-eps<=t[1]) & (t[1]<=1+eps)
            points = points[valid]
            prj_points = prj_points[valid]
            dist = dist[valid]
        return list(zip(abs(dist),points,prj_points))

### visualization functions
def plot_hull(hull, c=[1,1,1], ax=None):
    import matplotlib.pyplot as plt
    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
    pts = hull.points
    for s in hull.simplices:
        s = np.append(s, s[0])  # Here we cycle back to the first coordinate
        ax.plot(pts[s, 0], pts[s, 1], pts[s, 2], "-", color=c, marker = '.',linewidth=1.5)

def plot_bb(T,c,ax=None):
    import matplotlib.pyplot as plt
    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.set_xlim(-2,4)
        ax.set_ylim(-2,4)
        ax.set_zlim(-2,4)
    cube = np.array([  [-0.5, -0.5, -0.5, 1],
                       [ 0.5, -0.5, -0.5, 1],
                       [-0.5,  0.5, -0.5, 1],
                       [-0.5, -0.5,  0.5, 1],
                       [ 0.5,  0.5,  0.5, 1],
                       [-0.5,  0.5,  0.5, 1],
                       [ 0.5, -0.5,  0.5, 1],
                       [ 0.5,  0.5, -0.5, 1] ]).T  # 3D location of unit cube

    vertices = [[0, 1], [1, 7], [2, 7], [0, 2], [3, 6], [6, 4],
                [5, 4], [3, 5], [0, 3], [1, 6], [7, 4], [2, 5]] # vertices of a cube
    planes = [0,6]
    faces = [[0,2],[4,6],[0,4],[2,6],[1,5],[3,7]]

    points = T @ cube

    for f in faces:
        XYZ = np.array([points[:3,vertices[f[0]]],points[:3,vertices[f[1]]]])
        ax.plot_surface(XYZ[:,0,:],XYZ[:,1,:],XYZ[:,2,:], color=list(c)+[0.1], linewidth=0.5, edgecolor=list(c)+[0.8])
        ax.scatter(points[0],points[1],points[2],color=c)


from typing import Dict, List, Union, TypedDict

from ..interface import (
    load_watson_setting,
    load_whisper_setting,
    load_google_setting,
    EngineSettingInterface,
    PluginSettingsInterface
    )
from gailbot.core.utils.logger import makelogger
from gailbot.core.utils.general import write_toml
from gailbot.configs import default_setting_loader
logger = makelogger("engineObject")

class EngineSetObj():
    """
    store a single Engine setting

    """
    engine_setting: EngineSettingInterface   = None
    name: str     = None
    valid_interfaces = [load_whisper_setting, load_google_setting, load_watson_setting]

    def __init__(self, setting: Dict[str, str], name: str) -> None:
        """ initializing an engine

        Args:
            setting (Dict[str, str]): engine setting data stored in a dictionary
            name (str): the name of the engine setting
        """
        self.applied_in_profiles = set() # a set of profiles that applied this engine setting
        self.name = name
        logger.info("initialize the setting object")
        assert self._load_engine_setting(setting)
        self.data = setting
        self.engine = self.engine_setting.engine

    def get_init_kwargs(self):
        return self.engine_setting.get_init_kwargs()

    def get_transcribe_kwargs(self):
        return self.engine_setting.get_transcribe_kwargs()

    def get_name(self):
        """
        Accesses and returns the object's name
        """
        return self.name

    def change_name(self, name):
        """
        Changes the profile name to a given new name

        Args:
            name: name to change to
        """
        self.name = name

    def get_setting_dict(self) -> Dict[str, str]:
        """
        Accesses and returns the object's setting dict
        """
        return self.data

    def update_setting(self, setting: Dict[str, str]) -> bool:
        """
        Updates the settings to a given dictionary

        Args:
            setting: Dict[str, str]: new setting

        Returns:
            bool: True if successfully updated, false if not
        """
        logger.info(setting)
        self._load_engine_setting(setting)

        if self.engine_setting:
            self.data = setting
            return True
        else:
            return False

    def save_setting(self, output: str) -> bool:
        """
        Saves the settings to the output directory

        Args:
            output:str: output directory path

        Returns:
            bool: True if successfully saved, false if not
        """
        logger.info(self.data)
        try:
            write_toml(output, self.data)
        except Exception as e:
            logger.error(e, exc_info=e)
            return False
        else:
            return True

    def is_in_use(self):
        return len(self.applied_in_profiles) != 0

    def remove_applied_profile(self, profile_name):
        if profile_name in self.applied_in_profiles:
            self.applied_in_profiles.remove(profile_name)

    def add_applied_profile(self, profile_name):
        self.applied_in_profiles.add(profile_name)

    def _load_engine_setting(self, setting : Dict[str, str]) -> bool:
        """
        Loads the engine settings

        Args:
            setting : List[str]: settings to load

        Returns:
            bool: true if successfully loaded, false if not
        """
        logger.info("initialize the engine setting")
        for option in self.valid_interfaces:
            logger.info(option)
            set_obj = option(setting)
            if isinstance(set_obj, EngineSettingInterface):
                self.engine_setting = set_obj
                return True
        logger.info(f"setting {setting} cannot be loaded")
        return False


import requests
import json
import os
import re
from itertools import count

cnt = 0
stop = False
for offset in count(0):
    ch_url = 'http://www.podbbang.com/_m_api/podcasts/****/episodes'
    params = {
        'offset': offset,
    }

    html = requests.get(ch_url, params = params)
    print(html)

    result_dict = json.loads(html.text)

    tot_count = result_dict['summary']['total_count']
    print(tot_count)

    for r in result_dict['data']:
        cnt += 1
        print(cnt)

        print(r['id'])
        print(r['title'])
        print(r['enclosure']['url'])
        print(r['enclosure']['length'])
        print(r['duration'])
        print(r['published_at'])
        title = re.sub('[*\|<>?:"/]', '', r['title'])
        mp3_bin = requests.get(r['enclosure']['url']).content

        base_filename = os.path.basename(r['enclosure']['url'])
        filename = '{}_{}_{}'.format(r['id'], title, base_filename)
        print(filename)

        with open(filename, 'wb') as f:
            f.write(mp3_bin)

        print(cnt)

        if cnt >= tot_count:
            stop = True
            break

    if stop:
        break


from itertools import combinations
from operator import mul
import functools

def main():
    with open('advent2015_day24.txt') as file:
        weights = [int(x) for x in file.read().splitlines()]
        group_size = sum(weights) / 4
        for i in range(len(weights)):
            list = [functools.reduce(mul, c) for c in combinations(weights, i)
                    if sum(c) == group_size]
            if list:
                return min(list)

if __name__ == "__main__":
    print(main())

2015/Day_24 - sleigh balancing/advent2015_day24.py


import numpy as np
import random as rand
import matplotlib.pyplot as plt

arr = np.array([rand.randint(0,1000) for i in range(1000)])
arrnum = [i for i in range(1000)]
print("\nМассив:\n", arr,
      "\nМинимальное значение:\t",np.min(arr),
      "\nМаксимальное значение:\t",np.max(arr),
      "\nМедиана:\t\t",np.median(arr),
      "\nСтандартное отклонение:\t",np.std(arr),
      "\nСреднее значение:\t",np.mean(arr),
      "\nДисперсия:\t\t",np.var(arr))
plt.figure("Nice one",figsize = (12,6))

plt.subplot(111)
plt.bar(arrnum,arr)
plt.xlabel("Numbers")
plt.ylabel("Random values")

plt.show()


#!/usr/bin/env python3

from bcc import BPF
import ctypes as ct
import netaddr
import socket

# Constants
QTYPES = {'LOC': 29, '*': 255, 'IXFR': 251, 'UINFO': 100, 'NSEC3': 50, 'AAAA': 28, 'CNAME': 5, 'MINFO': 14, 'EID': 31, 'GPOS': 27, 'X25': 19, 'HINFO': 13, 'CAA': 257, 'NULL': 10, 'DNSKEY': 48, 'DS': 43, 'ISDN': 20, 'SOA': 6, 'RP': 17, 'UID': 101, 'TALINK': 58, 'TKEY': 249, 'PX': 26, 'NSAP-PTR': 23, 'TXT': 16, 'IPSECKEY': 45, 'DNAME': 39, 'MAILA': 254, 'AFSDB': 18, 'SSHFP': 44, 'NS': 2, 'PTR': 12, 'SPF': 99, 'TA': 32768, 'A': 1, 'NXT': 30, 'AXFR': 252, 'RKEY': 57, 'KEY': 25, 'NIMLOC': 32, 'A6': 38, 'TLSA': 52, 'MG': 8, 'HIP': 55, 'NSEC': 47, 'GID': 102, 'SRV': 33, 'DLV': 32769, 'NSEC3PARAM': 51, 'UNSPEC': 103, 'TSIG': 250, 'ATMA': 34, 'RRSIG': 46, 'OPT': 41, 'MD': 3, 'NAPTR': 35, 'MF': 4, 'MB': 7, 'DHCID': 49, 'MX': 15, 'MAILB': 253, 'CERT': 37, 'NINFO': 56, 'APL': 42, 'MR': 9, 'SIG': 24, 'WKS': 11, 'KX': 36, 'NSAP': 22, 'RT': 21, 'SINK': 40}
INV_QTYPES = {v: k for k, v in QTYPES.items()}
ACTIONS = {1 : 'DROP', 2 : 'TC'}

DROP_ACTION = 1
TC_ACTION = 2

# The interface on wich the filter will be attached
DEV = "eth0"

# The list of blocked IPv4, IPv6 and QNames
# IP format : (IPAddress, Action)
# CIDR format : (IPAddress/cidr, Action)
# QName format : (QName, QType, Action)
blocked_ipv4 = [("192.0.2.1", TC_ACTION)]
blocked_ipv6 = [("2001:db8::1", TC_ACTION)]
blocked_cidr4 = [("192.0.1.1/24", TC_ACTION)]
blocked_cidr6 = [("2001:db8::1/128", TC_ACTION)]
blocked_qnames = [("localhost", "A", DROP_ACTION), ("test.com", "*", TC_ACTION)]

# Main
xdp = BPF(src_file="xdp-filter.ebpf.src")

fn = xdp.load_func("xdp_dns_filter", BPF.XDP)
xdp.attach_xdp(DEV, fn, 0)

v4filter = xdp.get_table("v4filter")
v6filter = xdp.get_table("v6filter")
cidr4filter = xdp.get_table("cidr4filter")
cidr6filter = xdp.get_table("cidr6filter")
qnamefilter = xdp.get_table("qnamefilter")

for ip in blocked_ipv4:
  print(f"Blocking {ip}")
  key = v4filter.Key(int(netaddr.IPAddress(ip[0]).value))
  leaf = v4filter.Leaf()
  leaf.counter = 0
  leaf.action = ip[1]
  v4filter[key] = leaf

for ip in blocked_ipv6:
  print(f"Blocking {ip}")
  ipv6_int = int(netaddr.IPAddress(ip[0]).value)
  ipv6_bytes = bytearray([(ipv6_int & (255 << 8*(15-i))) >> (8*(15-i)) for i in range(16)])
  key = (ct.c_uint8 * 16).from_buffer(ipv6_bytes)
  leaf = v6filter.Leaf()
  leaf.counter = 0
  leaf.action = ip[1]
  v6filter[key] = leaf

for item in blocked_cidr4:
  print(f"Blocking {item}")
  key = cidr4filter.Key()
  network = netaddr.IPNetwork(item[0])
  key.cidr = network.prefixlen
  key.addr = socket.htonl(network.network.value)
  leaf = cidr4filter.Leaf()
  leaf.counter = 0
  leaf.action = item[1]
  cidr4filter[key] = leaf

for item in blocked_cidr6:
  print(f"Blocking {item}")
  key = cidr6filter.Key()
  network = netaddr.IPNetwork(item[0])
  key.cidr = network.prefixlen
  ipv6_int = int(network.network.value)
  ipv6_bytes = bytearray([(ipv6_int & (255 << 8*(15-i))) >> (8*(15-i)) for i in range(16)])
  key.addr.in6_u.u6_addr8 = (ct.c_uint8 * 16).from_buffer(ipv6_bytes)
  leaf = cidr6filter.Leaf()
  leaf.counter = 0
  leaf.action = item[1]
  cidr6filter[key] = leaf

for qname in blocked_qnames:
  print(f"Blocking {qname}")
  key = qnamefilter.Key()
  qn = bytearray()
  for sub in qname[0].split('.'):
    qn.append(len(sub))
    for ch in sub:
      qn.append(ord(ch))
  qn.extend((0,) * (255 - len(qn)))
  key.qname = (ct.c_ubyte * 255).from_buffer(qn)
  key.qtype = ct.c_uint16(QTYPES[qname[1]])
  leaf = qnamefilter.Leaf()
  leaf.counter = 0
  leaf.action = qname[2]
  qnamefilter[key] = leaf

print("Filter is ready")
try:
  xdp.trace_print()
except KeyboardInterrupt:
  pass

for item in v4filter.items():
  print(f"{str(netaddr.IPAddress(item[0].value))} ({ACTIONS[item[1].action]}): {item[1].counter}")
for item in v6filter.items():
  print(f"{str(socket.inet_ntop(socket.AF_INET6, item[0]))} ({ACTIONS[item[1].action]}): {item[1].counter}")
for item in cidr4filter.items():
  addr = netaddr.IPAddress(socket.ntohl(item[0].addr))
  print(f"{str(addr)}/{str(item[0].cidr)} ({ACTIONS[item[1].action]}): {item[1].counter}")
for item in cidr6filter.items():
  print(f"{str(socket.inet_ntop(socket.AF_INET6, item[0].addr))}/{str(item[0].cidr)} ({ACTIONS[item[1].action]}): {item[1].counter}")
for item in qnamefilter.items():
  print(f"{''.join(map(chr, item[0].qname)).strip()}/{INV_QTYPES[item[0].qtype]} ({ACTIONS[item[1].action]}): {item[1].counter}")

xdp.remove_xdp(DEV, 0)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Oct 13 15:59:34 2022

@author: craigrupp
"""

## Variance
# =============================================================================
# Sometimes we will deal with several samples, and in this case, it is vital to figure out if there is any relationship between two samples.
#
# We will start with discovering two independent samples; let's assume that they are distributed normally, and we know that we are working with the sample, not the whole population.
# Our task is to figure out if the sample means values are equal or not to conclude the sample's difference.
#
# Firstly we should create the null hypothesis stating that two different mean values are equal.
# But if we reject this hypothesis, we move to the second one that says that our mean values are different and, consequently, the samples do not belong to one population.
#
# Null Hypothesis : H(0) - Two different means values are equal (μ1 = μ2)
# Alternative : H(1) - Two Samples' mean are different (μ1 != μ2)
# =============================================================================
# =============================================================================
# Student's T-test : Small Dist
# sem :  standard error of the mean
# T = μ1 - μ2 / sem
# =============================================================================

## T-Test Conditions
# Samples are distributed normally.
# Homogeneity of variance : It is when we assume the equality or similarity of variances of two samples. (Levene Test - PValue Hello!)

## Levene's test (P-Value) (method that you will use not only in the case with normal distribution; it is universal and fits the situation when the distribution is not normal or has outliers!)
# =============================================================================
# If it is a test, we should state the null hypothesis. To do it, firstly, figure out the p-value.
# There is nothing complicated here; it is just a value that we use to reject or not the null hypothesis; usually, the p-value equals 0.05.
# But in the test, we will receive another p-value, let's call it a p-test.
# If p-test < 0.05 (p-test < p-value), we reject the null hypothesis and reject the hypothesis that variances of two samples are equal; otherwise, we accept it, and variances are equal.
# =============================================================================

## Task
# =============================================================================
# Assume that recently you've created two apps and want to know if they are equally successful.
# You have two samples about the mean amount of users that download an app every day for one month.
#
# sample_1 describes the first app with 87 observations, a mean value of 35, and a standard deviation of 2.78.
# sample_2 represents the second app with 79 observations, a mean value of 39.1, and a standard deviation of 2.99.
# Your task here is to check for the homogeneity of variances
# =============================================================================
# Import levene object
from scipy.stats import levene
# Import norm object
from scipy.stats import norm
import numpy as np

np.random.seed(0)

sample_1 = norm.rvs(size = 87, loc = 35, scale = 2.78)
sample_2 = norm.rvs(size = 79, loc = 39.1, scale = 2.99)

# Define the p-value
p_value = 0.05
# Conduct the Levene`s test
stat, p_test = levene(sample_1, sample_2)

print("Our p-value is", p_test)
# Check if we need to reject the null hypothesis
if p_test <  p_value:
    print("The variances of two samples are different")
else:
    print("The variances of two samples are equal")

# Our p-value is 0.462614712112142
# The variances of two samples are equal

## Equal Variance :  (Test Conditions Met with Homogeneity of variance above and sample being known to be distributed normally)
# =============================================================================
# The variances of two samples are equal this phrase is the result from the previous chapter! It means that the variances of the two samples are very similar,
# and the samples do not vary from each other at all. So, if the variance is equal and the samples are distributed normally, we can conduct the test because
# all requirements were met!
# =============================================================================
# =============================================================================
# ## Pooled Variance
# # If we know that the variances are equal, we can estimate the correlation between pooled data set (two samples) and the average of samples to conclude if they are identical.
# ## Equation
# # Pooled Variance = (sample1_var * deg_freedmom_sample1) + (sample2  * deg_freedmom_sample2) / deg_freedmom_sample1  + deg_freedmom_sample2
# =============================================================================

# =============================================================================
# ## Standard error : is a mathematical tool used in statistics to measure variability.
# # It enables one to arrive at an estimation of what the standard deviation of a given sample is.
# # It is commonly known by its abbreviated form – SE. Standard error is used to estimate the efficiency, accuracy, and consistency of a sample
# ## Equation : se = np.sqrt(pooled_var * (1 / len_1 + 1 / len_2))
# =============================================================================

len_1 = len(sample_1)
# Calculate the length of the second sample
len_2 = len(sample_2)

# Calculate the degrees of freedom for the first sample
df_1 = len_1 - 1
# Calculate the degrees of freedom for the second sample
df_2 = len_2 - 1

# Calculate the mean for the first sample
sample_1_mean = np.mean(sample_1)
# Calculate the mean for the second sample
sample_2_mean = np.mean(sample_2)

# Calculate the variance for the first sample
sample_1_var = np.var(sample_1)
# Calculate the variance for the second sample
sample_2_var = np.var(sample_2)

# Calculate the pooled variance using two samples who passed test
pooled_var = (sample_1_var * df_1 ) + (sample_2_var * df_2) / (df_1 + df_2)
# Calculate the standard error
se = np.sqrt(pooled_var * (1 / len_1 + 1 / len_2))

print("The mean of the first sample is", sample_1_mean)
print("The mean of the second sample is", sample_2_mean)
print("The variance of the first sample is", sample_1_var)
print("The variance of the second sample is", sample_2_var)
print("The pooled variance is", pooled_var)
print("The standard error is", se)

# The mean of the first sample is 35.02507220040284
# The mean of the second sample is 39.776864153766326
# The variance of the first sample is 8.260950582885165
# The variance of the second sample is 8.726647059072187
# The pooled variance is 714.5922286074391
# The standard error is 4.154416327872356

## Two-Tailed T-Test
# =============================================================================
# It is the test that helps us define if there is any difference in samples(increasing or decreasing in values).
# Another wording allows us to determine if samples are equal or not.
#
# (With Significance Level at 0.05)
# We should divide the significance level by two because it is a two-tailed test, and if our value falls into 2.5% data from the right,
# it means that the mean of the first sample is greater than the second; otherwise, if our value falls into 2.5% data from the left,
# it means that the mean of the first sample is less than the second. If, in general, we claim that a = 0.5, the confidence level will be 1 - a/2 = 0.975.
# =============================================================================
## Conduct T-Test
# =============================================================================
# 1. Find the value of t-students statistics (value of t), but it is better to use a module of the differences of the mean values, let it be t variable.
#     t = μ1 - μ2 / SE (Standard Error)
# 2. Find the critical value for the two-tailed test,
#     t-critical: t_critical = stats.t.ppf(q=0.975, df=df)
# here we conduct a test with the a value equal to 0.05, and put a sum of degrees of freedom of the explored samples.
# 3. Find the new p-value to which we will compare p-critical in our t-test :
#         p_value_two_tailed = 2*(1 - stats.t.cdf(x = t, df = df))
# as the first argument, we put the value of t-statistics, and as the second one, we put the sum of degrees of freedom for two samples.
# 4. Check if the p_value_two_tailed is more significant than a value of 0.05. We support the null hypothesis if this condition is proper; otherwise, we reject it.
# =============================================================================
# bring in stats library
import scipy.stats as st
# Calculate t value
t = abs(sample_1_mean - sample_2_mean) / se
# Calculate critical value : a level 0.05
t_critical = st.t.ppf(q = 0.975, df = df_1 + df_2)
# Calculate p-value
p_value_two_tailed = 2 * (1 - st.t.cdf(x = t, df = df_1 + df_2))
# Check the null hypothesis
if p_value_two_tailed > 0.05:
    print("We support the null hypothesis, the mean values are equal")
# Check the alternative hypothesis
else:
    print("We reject the null hypothesis, the mean values are different")
print("t-statistics value equal to", t)
print("t-critical value equal to", t_critical)

# We reject the null hypothesis, the mean values are different
# t-statistics value equal to 10.498231324238853
# t-critical value equal to 1.974534575852265

## Python Conduct Sample T-Test
# =============================================================================
# Python formula for the sample t-test
#
# st.ttest_ind(a = sample_1, b = sample_2, equal_var=True), it returns two values: the first one is t-statistics and the second one is the p-value.
# The arguments are pretty simple; the first and second corresponding to the discovered samples, and the third define homogeneity of variance.
# =============================================================================

## Task
# As we check the homogeneity of variance and know that the samples are normally distributed, we can boldly apply the t-test to our samples.

# Apply t-test
stats, p_test = st.ttest_ind(a = sample_1,
                          b = sample_2,
                          equal_var = True)

# Check if we should support or nit the null hypothesis
if p_test > 0.05:
    print("We support the null hypothesis, the mean values are equal")
else:
    print("We reject the null hypothesis, the mean values are different")

# We reject the null hypothesis, the mean values are different

## Paired T-Test
# =============================================================================
# You finished the first test and figured out that the mean value of users that download the app is different!
# So, now it is time to dive into a fascinating test that looks almost the same but is just a little bit complicated.
# =============================================================================

# Scenario
# =============================================================================
# Let's assume that you have a successful app and want to figure out if the number of users that download the app will statistically increase
# during one month if you change the logo. We should have statistical proof that the mean value has decreased; we can not just compare two mean values and conclude.
# So, you have two samples: one before changes and another after. Such samples are called dependent and require paired t-test.
# =============================================================================

# Conditions for Test
# =============================================================================
# Samples should be normally distributed.
# Samples should not contain outliers.
# Samples should have equal lengths.
# =============================================================================

# Conduct - stats.ttest_rel(sample_a, sample_b)
# =============================================================================
# st.ttest_rel(a, b) is similar to the function for independent samples, but here we do not obligate to check the homogeneity of variance;
# we just put two samples that should be discovered. In the output, you will receive two values the first one is t-statistics and the second one is the p-value.
# The condition for supporting or rejecting the null hypothesis is the same for all tests.
# =============================================================================
np.random.seed(0)

sample_1_paired = norm.rvs(size = 104, loc = 91, scale = 2.78)
sample_2_paired = norm.rvs(size = 104, loc = 92, scale = 2.78)

# Conduct paired t-test
stats, p_test = st.ttest_rel(sample_1_paired, sample_2_paired)

# Check the null hypothesis
if p_test > 0.05:
    print("We support the null hypothesis, the mean values are equal")
# Check the alternative hypothesis
else:
    print("We reject the null hypothesis, the mean values are different")

## Mann-Whitney U test
# =============================================================================
# We have already learned how to conduct the t-test, and we have recognized that this test requires a lot!
# Our samples should be normally distributed, and the variances should be equal... But what should we do if we want to compare two independent samples but do not want to change them?
# =============================================================================

## Scenario (See Image in Ghub Repo for Assocation - Ranging Concept - Big to Small for each sample (high to low or vice versa (doesn't matter)) then count sum of ranged values for each group)
# =============================================================================
# The main principle is to divide your data into groups and range it. Look at the picture where we can observe houses from the first city(purple houses) and houses from the second city(yellow houses), and they vary because of their size:
#
# 1) In the first part of the picture, you can see six different houses of different sizes.
# 2) In the second part of the picture, we ranged all houses; the first one (the biggest) has the range 1, the second one (smaller than the biggest one) has the range 2 and the smallest one here has the range 6.
# 3) Then we divide our houses by cities and count the sum of ranges, and if the sum of ranges for different cities varies significantly, we can assume that the sizes of houses vary too and vice versa.
# =============================================================================

## Next Step - U-test
# =============================================================================
# The first step in conducting a u-test is to determine if the shape of distributions is equal;
# they have equal amounts of peaks and approximately similar types. If they are equal, you should compare medians of two samples; otherwise, mean values.
# =============================================================================

## Task
# =============================================================================
# We will discover the prices of cars Ford Focus and Ford Fiesta; our task is to determine if the mean prices of the two models are equal.
# If not, we will discover which model is more expensive.
# Firstly let's compare the distributions' shapes
# =============================================================================
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
# Read csv file about Ford Fiesta
data_fiesta = pd.read_csv('https://codefinity-content-media.s3.eu-west-1.amazonaws.com/a849660e-ddfa-4033-80a6-94a1b7772e23/section_1/Fiesta')
# Read csv file about Ford Focus
data_focus = pd.read_csv('https://codefinity-content-media.s3.eu-west-1.amazonaws.com/a849660e-ddfa-4033-80a6-94a1b7772e23/section_1/Focus')

# Build a histplot for Ford Fiesta
sns.histplot(data = data_fiesta,
            x = 'price',
            color = 'b',
            kde = True,
            label = 'Fiesta')
# Build a histplot for Ford Focus
sns.histplot(data = data_focus,
            x = 'price',
            color = 'y',
            kde = True,
            label = 'Focus')

plt.legend()
plt.show()

## Task Cont : Check If the Data Has Outliers - Box Plots!
# =============================================================================
# As you may recognize, the shapes of our data are not equal; one distribution even has two peaks, so the right way here is to compare mean values.
#
# The next step in our test is not obligatory; we should be sure that u-test is appropriate here, so as you remember, this test can deal with samples with outliers.
# The best way to check them is to create boxplots! So, your task in this chapter is just to ascertain that our data are not so beautiful
# =============================================================================
# Create DataFrame
price = pd.DataFrame({'Focus': data_focus['price'], 'Fiesta' : data_fiesta['price']})
print(price.tail())
print(price['Focus'].isna().sum())
print(price['Fiesta'].isna().sum())

# Build a boxplot using Dataframe with dataset price columns
price.boxplot(column=['Focus', 'Fiesta'])
plt.ylabel('Price')
plt.show()

## Task Cont : Compare Mean - st.mannwhitneyu
# =============================================================================
# We've prooved our data full of outliers, so we can conduct U-test for the prices of Ford cars.
#
# As usual, let's figure out the hypothesis:
#
# The null hypothesis: the mean prices for two cars are equal;
# The alternative hypothesis: are the mean prices for the two cars vary.
# The condition for supporting the null hypothesis is the same: if the p-value received in the test is greater than the defined p-value(0.05),
# we support the null hypothesis; otherwise, we reject it.
#
# Function for calculating U-test:
# statics, p_value = st.mannwhitneyu(x = sample_1, y = sample_2, alternative = 'two-sided')
#
# x = sample_1 the first sample,
# y = sample_2 the second sample,
# alternative = 'two-sided' means that we are providing a two-tailed test here to figure out are the means different (now, it is no matter if one mean is bigger than another).
# As a result, we will receive the u-statistics variable the p_value, which should be compared.
# =============================================================================
# Conduct U-test
statics, p_value = st.mannwhitneyu(x = data_focus['price'],
                         y = data_fiesta['price'],
                         alternative = 'two-sided')
# Print the results
print(statics, p_value)

# Check the null hypothesis
if p_value > 0.05:
    print("We support the null hypothesis, the mean values are equal")
else:
    print("We reject the null hypothesis, the mean values are different")

# 21463134.5 0.0
# We reject the null hypothesis, the mean values are different

## Deal with One-Tailed Test - Which of Two samples has a greater mean value
# =============================================================================
# We know that mean values are different, but it is time to check which model, on average, is more expensive.
#
# It is time for you to dive deeper into the one-tailed test. It differs from the two-tailed test because this one helps
# us find out which of the two samples has a greater mean value
# =============================================================================
## What is the one-tailed test?
# =============================================================================
# Imagine that we have two samples and, consequently, two hypotheses.
#
# # First Case
# The null hypothesis: the first sample is greater than or equal to the second sample.
# The alternative hypothesis, the second sample is less than the first sample.
#
# In the first picture, you can see the right one-tailed test that helps us figure out if the first sample is greater than the second;
# it means that our data should fall to the left of the red line; otherwise, we reject the null hypothesis.
#
# # In the second case, let's assume:
#
# The null hypothesis: the first sample is less than the second sample.
# The alternative hypothesis, the second sample is greater than or equal to the first sample.
#
# In the second picture, you can see the left one-tailed test that helps us figure out if the first sample is less than the second,
# which means that our data should fall to the right of the red line; otherwise, we reject the null hypothesis.
# =============================================================================
## Functions
# =============================================================================
# statics, p_value = st.mannwhitneyu(x = sample_1, y = sample_2, alternative = 'less')
# - check if the sample_1 mean is less than the sample_2 mean.(alternative hypothesis: sample_1 is less)
#
# statics, p_value = st.mannwhitneyu(x = sample_1, y = sample_2, alternative = 'greater')
# - check if the sample_1 mean is greater thsn the sample_2 mean (alternative hypothesis: sample_1 is greater).
# =============================================================================

## Task
# =============================================================================
# Let's assume the hypotheses:
#
# The null hypothesis: the Ford Fiesta model cheaper than the Ford Focus model.
# The alternative hypothesis: the Ford Fiesta model more expensive or equal to the Ford Focus model.
# Follow the algorithm to reject one of them:
# =============================================================================
# Conduct U-test
statics, p_value = st.mannwhitneyu(x = data_fiesta['price'],
                      y = data_focus['price'],
                      alternative = 'greater')
print(statics, p_value)

# Check the hypothesis
if p_value > 0.05:
    print("We support the null hypothesis, the mean price of Ford Fiesta is less than the Ford Focus")
else:
    print("We reject the null hypothesis, the mean price of Ford Fiesta is greater than or equal to the Ford Focus")

## 8620381.5 1.0
## We support the null hypothesis, the mean price of Ford Fiesta is less than the Ford Focus

## Another Hypotheses
# =============================================================================
# Let's assume another hypotheses, to test out u-test in Python:
#
# The null hypothesis: the Ford Fiesta model more expensive than or equal to the Ford Focus model.
# The alternative hypothesis: the Ford Fiesta model less expensive than the Ford Focus model.
# Follow the algorithm to reject one of them:
# =============================================================================

data_fiesta = pd.read_csv('https://codefinity-content-media.s3.eu-west-1.amazonaws.com/a849660e-ddfa-4033-80a6-94a1b7772e23/section_1/Fiesta')
data_focus = pd.read_csv('https://codefinity-content-media.s3.eu-west-1.amazonaws.com/a849660e-ddfa-4033-80a6-94a1b7772e23/section_1/Focus')

# Conduct U-test
statics, p_value = st.mannwhitneyu(x = data_fiesta['price'],
                                 y = data_focus['price'],
                                 alternative = 'less')
print(statics, p_value)

# Check the hypothesis
if p_value > 0.05:
    print("We support the null hypothesis, the mean price of Ford Fiesta is greater than or equal to the Ford Focus")
else:
    print("We reject the null hypothesis, the mean price of Ford Fiesta is less than the Ford Focus")

# Check the mean values
print(np.mean(data_fiesta['price']), np.mean(data_focus['price']))

# =============================================================================
# 8620381.5 0.0
# We reject the null hypothesis, the mean price of Ford Fiesta is less than the Ford Focus
# 10196.298002135123 13185.88295553618
# =============================================================================

## Seems The Null Hypothesis is essentially just the inverse of our alternative argument based on the positioning of our two samples
# =============================================================================
# first hypotheses (lines 415 - 417) : Null Hypothesis confirmed for fiesta price being less w/alternative set to greater
# second hypothes (lines 446 - 448) : Null Hypothesis rejected for fiesta being greater or equal to ford focus w/alternative set to less
# =============================================================================


# pylint: skip-file
from openpyxl import Workbook
from openpyxl.styles import Border, Side, Alignment
from datetime import datetime
from openpyxl.utils import get_column_letter

def create_formula(column_shift:int, player_id:int, number_of_rounds:int, number_of_tables:int, number_of_players:int) -> str:
    # Create the formula for the total tichu or victory points of a player
    formula = "="
    first_occurence = False
    for table_id in range(number_of_tables):
        # Create the formula for each table
        for team in range(2):
            # Repeat it twice (for the two teams of each table)
            for seat in range(2):
                # Repeat it twice (for the two seats of each team)
                if first_occurence:
                    # Add a "+" between the formulas on every occurence except the first one
                    formula += "+"
                first_occurence = True
                # Concatenate the formulas for each of the two column per table
                formula += f"SUMIF(${get_column_letter((table_id*9)+3+seat+(team*4))}{number_of_players+6}:{get_column_letter((table_id*9)+3+seat+(team*4))}{number_of_players+number_of_rounds+5},$B{player_id+1},{get_column_letter((table_id*9)+5+column_shift+(team*4))}{number_of_players+6}:{get_column_letter((table_id*9)+5+column_shift+(team*4))}{number_of_players+number_of_rounds+5})"
    return formula

def create_player_table(players:list, number_of_rounds:int, number_of_tables:int, number_of_players:int, ws:Workbook) -> None:
    # Create the player table to insert the names of the players
    ws.cell(row=1, column=1).value = "Spieler"
    ws.cell(row=1, column=2).value = "Name"
    ws.cell(row=1, column=3).value = "TP gesamt"
    ws.cell(row=1, column=4).value = "SP gesamt"
    for player_id in players:
        # Add the names of the players
        ws.cell(row=player_id+1, column=1).value = player_id
        ws.cell(row=player_id+1, column=2).value = player_id
        # Add the formulas for the total tichu and victory points
        ws.cell(row=player_id+1, column=3).value = create_formula(0, player_id, number_of_rounds, number_of_tables, number_of_players)
        ws.cell(row=player_id+1, column=4).value = create_formula(1, player_id, number_of_rounds, number_of_tables, number_of_players)

def create_headers(number_of_players:int, number_of_tables:int, ws:Workbook) -> None:
    # Create the headers for the tournament plan
    offset_y = number_of_players + 4
    ws.cell(row=offset_y+1, column=1).value = "Runde"
    for table_id in range(number_of_tables):
        # Add headers for each table
        offset_x = (table_id*9) + 3
        table = ws.cell(row=offset_y, column=offset_x)
        table.value = f"Tisch {table_id+1}"
        ws.merge_cells(start_row=offset_y, start_column=offset_x, end_row=offset_y, end_column=offset_x+7)
        table.alignment = Alignment(horizontal='center')
        for team_id in range(1,3):
            # Add headers for team 1 and team 2
            team_offset = offset_x+(team_id*4)-4
            team = ws.cell(row=offset_y+1, column=team_offset)
            ws.merge_cells(start_row=offset_y+1, start_column=team_offset, end_row=offset_y+1, end_column=team_offset+1)
            team.value = f"Team {team_id}"
            team.alignment = Alignment(horizontal='center')
            # Add headers for tichu points and victory points
            tichu_points = ws.cell(row=offset_y+1, column=team_offset+2)
            victory_points = ws.cell(row=offset_y+1, column=team_offset+3)
            tichu_points.value = "TP"
            victory_points.value = "SP"

def create_tournament_table(tournament_plan:list, number_of_players:int, ws:Workbook) -> None:
    # Create the tournament table with the tournament plan
    for round_id,round in enumerate(tournament_plan):
        offset_y = round_id + number_of_players + 6
        ws.cell(row=offset_y, column=1).value = round_id+1
        for table_id,table in enumerate(round):
            offset_x = (table_id*9) + 2
            for i in range(4):
                # Insert the players for each team
                player = ws.cell(row=offset_y, column=(i+1)+offset_x)
                player.value = f"=B{table[i]+1}"
                if i == 0:
                    player.border = Border(left=Side(border_style="thin",color='00000000'))
                if i == 1 or i == 3:
                    # Insert cells for tichu points and victory points for each team
                    tichu_points = ws.cell(row=offset_y, column=(i+1)+offset_x+1)
                    victory_points = ws.cell(row=offset_y, column=(i+1)+offset_x+2)
                    tichu_points.value = 0
                    victory_points.value = 0
                    victory_points.border = Border(right=Side(border_style="thin",color='00000000'))
                    offset_x += 2

def create_excel(players:list, tournament_plan:list, number_of_players:int, number_of_tables:int, number_of_rounds:int) -> None:
    # Create the excel file with the tournament plan
    today = datetime.today().strftime('%d_%m_%Y')
    wb = Workbook()
    ws = wb.active
    create_player_table(players, number_of_rounds, number_of_tables, number_of_players, ws)
    create_headers(number_of_players, number_of_tables, ws)
    create_tournament_table(tournament_plan, number_of_players, ws)
    wb.save(f'tichu_tournament_{today}.xlsx')

if __name__ == '__main__':
    players = [x for x in range(1, 9)]
    number_of_players = len(players)
    number_of_tables = number_of_players // 4
    number_of_rounds = 6
    tournament_plan =  [[[5, 7, 2, 6], [4, 1, 8, 3]],
                        [[3, 5, 8, 4], [1, 6, 2, 7]],
                        [[2, 1, 6, 5], [8, 7, 4, 3]],
                        [[5, 1, 6, 4], [3, 7, 8, 2]],
                        [[1, 8, 3, 2], [6, 7, 4, 5]],
                        [[7, 1, 2, 4], [5, 8, 3, 6]]]
    create_excel(players, tournament_plan, number_of_players, number_of_tables, number_of_rounds)


import sys
import time
from collections import defaultdict

import dill
from googletrans import Translator

from display_occurrences import load_dicts

SOURCE_LANG = 'sv'
TARGET_LANG = 'en'
MAX_LEN = 4970

def translate(source_folder, target_file):
    translator = Translator()

    complete_dictionary = load_dicts(source_folder)
    translated_dictionary = defaultdict(lambda: defaultdict(int))

    for pos in complete_dictionary:
        source_batch = []
        occurrences_batch = []
        source_bulk = ''
        occurrences = []
        translated_batch = []

        sorted_words = sorted(complete_dictionary[pos].items(), key=lambda x: x[1], reverse=True)
        for word, occurrence in sorted_words:
            source_bulk += word + '\n'
            occurrences.append(occurrence)

            if len(source_bulk) > MAX_LEN:
                source_batch.append(source_bulk)
                source_bulk = ''
                occurrences_batch.append(occurrences)
                occurrences = []
        source_batch.append(source_bulk)
        occurrences_batch.append(occurrences)

        for index, source_text in enumerate(source_batch):
            print(pos, index)
            delays = [0, 5, 10, 20, 30]
            for delay in delays:
                try:
                    if delay > 0:
                        print(f'Waiting {delay}s...')
                    time.sleep(delay)
                    translated_batch.append(translator.translate(source_text, src=SOURCE_LANG, dest=TARGET_LANG))
                    break
                except Exception:
                    continue

        for batch_index, translated_text in enumerate(translated_batch):
            separated_source_words = source_batch[batch_index].split('\n')
            separated_translated_words = translated_text.text.split('\n')
            for word_index, translated_word in enumerate(separated_translated_words):
                translated_dictionary[pos][(separated_source_words[word_index], translated_word.lower())] = \
                    occurrences_batch[batch_index][word_index]

    with open(target_file, 'wb') as f:
        dill.dump(translated_dictionary, f)

if __name__ == '__main__':
    if len(sys.argv) == 1:
        raise ValueError('Following arguments mandatory: source_folder, target_file')
    translate(sys.argv[1], sys.argv[2])