import os
import argparse
import numpy as np
import pandas as pd
from sklearn.model_selection import StratifiedKFold

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Arguments')
    parser.add_argument('--data_csv_path', type=str,
                        help='path to the csv file containing the cell expression data of all cases along with their labels')
    parser.add_argument('--output_dir', type=str, help='directory to processed data files')

    args = parser.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    # data should be downloaded from the link given in the paper.
    data = pd.read_csv(os.path.join(args.data_csv_path))

    # Merge Cytotoxic CD8 cells to CD8 cells
    data.loc[data['Annotation'] == 'Cytotoxic CD8', 'Annotation'] = 'CD8'
    data.reset_index(inplace=True)

    # Extracting and Saving Marker Names from the data CSV File
    non_marker_cols = ['patientID', 'centroidX', 'centroidY', 'cellLabel', 'cellSize', 'Annotation', 'pointNum',
                       'tissue_block', 'region', 'index']
    marker_names = [col_name for col_name in data.columns.tolist() if col_name not in non_marker_cols]
    marker_names.sort()
    marker_df = pd.DataFrame(marker_names, columns=['marker_name'])
    marker_df.to_csv(os.path.join(args.output_dir, 'marker_names.csv'), index=False)

    # Creating Class ID and Name Mapping for Unique Classes in a Dataset

    class_dict = {'class_id': [], 'class_name': [], 'sample_count': []}
    class_mapping = {}
    for i, class_name in enumerate(np.unique(data['Annotation'].tolist())):
        class_dict['class_id'].append(i)
        class_dict['class_name'].append(class_name)
        class_dict['sample_count'].append(np.sum(data['Annotation'] == class_name))
        class_mapping[class_name] = i
    pd.DataFrame(class_dict).to_csv(os.path.join(args.output_dir, 'class_names.csv'), index=False)

    # Count number cells per patient per class
    class_names = np.unique(data['Annotation']).tolist()
    case_names = np.unique(data['patientID']).tolist()
    case_names = [case_name for case_name in case_names if 'ctrl' not in case_name]
    df = pd.DataFrame(np.zeros((len(case_names), len(class_names))), columns=class_names, index=case_names)
    for case_name, row in df.iterrows():
        for class_name in class_names:
            temp_df = data[data['patientID'] == case_name]
            temp_df = temp_df[temp_df['Annotation'] == class_name]
            df.loc[case_name, class_name] = temp_df.shape[0]
    df['Total'] = np.sum(df.to_numpy(), axis=1)
    df.to_csv(os.path.join(args.output_dir, 'class_count.csv'))

    # Assigning fold id to each patient based on number of cells
    df = df.sort_values(['Total'])
    case_to_fold_mapping = {'case_name': df.index.tolist(), 'fold_id': [-1] * df.shape[0]}

    i, j = 0, df.shape[0] - 1
    fid = 0
    fold_count = 5
    cell_per_class = df.to_numpy()
    cell_per_fold = np.zeros((fold_count, df.shape[1]))
    while i < j:
        case_to_fold_mapping['fold_id'][i] = fid % fold_count
        case_to_fold_mapping['fold_id'][j] = fid % fold_count
        cell_per_fold[fid % fold_count, :] += cell_per_class[i, :] + cell_per_class[j, :]
        fid += 1
        i += 1
        j -= 1
    cell_per_fold = pd.DataFrame(cell_per_fold, columns=df.columns)
    cell_per_fold.to_csv(os.path.join(args.output_dir, 'cell_count_per_fold.csv'))
    case_to_fold_mapping = pd.DataFrame(case_to_fold_mapping)
    case_to_fold_mapping.to_csv(os.path.join(args.output_dir, 'case_to_fold_mapping.csv'))

    # Generating training and validation splits using fold information
    features_cols = []
    features_cols.extend(marker_names)
    features_cols.append('cellSize')
    features_cols.append('patientID')
    features = data[features_cols]
    labels = np.array([class_mapping[class_name] for class_name in data['Annotation'].tolist()])
    features['cell_label'] = labels

    for i in range(fold_count):
        train_patient_list = case_to_fold_mapping[case_to_fold_mapping['fold_id'] != i]['case_name'].tolist()
        test_patient_list = case_to_fold_mapping[case_to_fold_mapping['fold_id'] == i]['case_name'].tolist()

        train_split = features[features['patientID'].isin(train_patient_list)]
        test_split = features[features['patientID'].isin(test_patient_list)]
        train_split = train_split.drop(columns=['patientID'])
        test_split = test_split.drop(columns=['patientID'])

        print(i, train_split.shape, test_split.shape)

        os.makedirs(os.path.join(args.output_dir, 'splits', 'fold_%d' % i), exist_ok=True)
        train_split.to_csv(os.path.join(args.output_dir, 'splits', 'fold_%d' % i, 'train.csv'), index=False)
        test_split.to_csv(os.path.join(args.output_dir, 'splits', 'fold_%d' % i, 'test.csv'), index=False)

    # Generating training, validation and test splits using fold information
    for fid in range(5):
        os.makedirs(os.path.join(args.output_dir, 'splits', f"fold_{fid}"), exist_ok=True)
        train_csv = pd.read_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'train.csv'))
        features = train_csv.to_numpy()[:, :-1]
        labels = train_csv.to_numpy()[:, -1]
        skf = StratifiedKFold(n_splits=5, random_state=7325111, shuffle=True)
        for i, (train_index, valid_index) in enumerate(skf.split(features, labels)):
            train = train_csv.iloc[train_index]
            valid = train_csv.iloc[valid_index]
            print(f"Fold {fid}: {train.shape[0]}, {valid.shape[0]}")
            train.to_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'train.csv'), index=False)
            valid.to_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'valid.csv'), index=False)
            break

    # Generating subset of training sets in each fold for training the model with different number of training samples
    for fid in range(5):
        train_csv = pd.read_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'train.csv'))
        valid_csv = pd.read_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'valid.csv'))
        test_csv = pd.read_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'test.csv'))
        features = train_csv.to_numpy()[:, :-1]
        labels = train_csv.to_numpy()[:, -1]
        train_index = []
        skf = StratifiedKFold(n_splits=20, random_state=7325111, shuffle=True)
        for i, (_, train_index_) in enumerate(skf.split(features, labels)):
            train_index.extend(train_index_)
            if i in [0, 1, 4, 9, 14]:
                train = train_csv.iloc[train_index]
                print(
                    f"Fold {fid}: {(i + 1) * 5}%: {len(train_index) / features.shape[0]}, {train.shape[0] / features.shape[0]}")
                os.makedirs(os.path.join(args.output_dir, 'splits_%02d' % ((i + 1) * 5), f"fold_{fid}"), exist_ok=True)
                train.to_csv(os.path.join(args.output_dir, 'splits_%02d' % ((i + 1) * 5), f"fold_{fid}", 'train.csv'), index=False)
                valid_csv.to_csv(os.path.join(args.output_dir, 'splits_%02d' % ((i + 1) * 5), f"fold_{fid}", 'valid.csv'), index=False)
                test_csv.to_csv(os.path.join(args.output_dir, 'splits_%02d' % ((i + 1) * 5), f"fold_{fid}", 'test.csv'), index=False)

    # Generating train and valid sets with subset markers and classes
    selected_cols = ["CD11c", "CD163", "CD20", "CD4", "CD56", "CD68", "CD8", "CD3", "Pax-5", "cellSize", 'cell_label']

    pd.DataFrame(selected_cols, columns=['col_names']).to_csv(os.path.join(args.output_dir, 'selected_cols.csv'), index=False)
    selected_classes = ["CD4", "CD8", "DC", "M1", "M2", "NK", "Tumor", "B"]
    selected_labels = [class_mapping[class_name] for class_name in selected_classes]
    selected_labels_mapping = {class_mapping[class_name]: i for i, class_name in enumerate(selected_classes)}
    print(selected_labels_mapping)
    selected_labels_mapping[class_mapping['B']] = selected_classes.index('Tumor')
    print(selected_labels_mapping)
    class_dict = {'class_id': [], 'class_name': []}
    class_mapping = {}
    for i, class_name in enumerate(selected_classes):
        if class_name == 'B':
            continue
        class_dict['class_id'].append(i)
        class_dict['class_name'].append(class_name)
        class_mapping[class_name] = i
    pd.DataFrame(class_dict).to_csv(os.path.join(args.output_dir, 'selected_class_names.csv'), index=False)

    for fid in range(5):
        train = pd.read_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'train.csv'))
        valid = pd.read_csv(os.path.join(args.output_dir, 'splits', f"fold_{fid}", 'valid.csv'))

        train = train[selected_cols]
        valid = valid[selected_cols]

        train = train[train['cell_label'].isin(selected_labels)]
        valid = valid[valid['cell_label'].isin(selected_labels)]

        train['cell_label'] = train['cell_label'].map(selected_labels_mapping)
        valid['cell_label'] = valid['cell_label'].map(selected_labels_mapping)

        os.makedirs(os.path.join(args.output_dir, 'splits_marker_subset', f"fold_{fid}"), exist_ok=True)
        train.to_csv(os.path.join(args.output_dir, 'splits_marker_subset', f"fold_{fid}", 'train.csv'), index=False)
        valid.to_csv(os.path.join(args.output_dir, 'splits_marker_subset', f"fold_{fid}", 'valid.csv'), index=False)


import os
import sys

import progressbar as pb

from vinfraclient import exceptions
from vinfraclient import utils
from vinfraclient.argtypes import parse_list_options
from vinfraclient.cmd import base

def _image_arg(parser):
    parser.add_argument(
        "image",
        metavar="<image>",
        help="Image ID or name"
    )

def _common_set_options(parser):
    parser.add_argument(
        "--min-disk",
        type=int,
        metavar="<size-gb>",
        help=" Minimum disk size required to boot from image, in gigabytes"
    )
    parser.add_argument(
        "--min-ram",
        type=int,
        metavar="<size-mb>",
        help="Minimum RAM size required to boot from image, in megabytes"
    )
    parser.add_argument(
        "--os-distro",
        metavar="<os-distro>",
        help="OS distribution. To list available distributions, "
             "run 'vinfra service compute show'."
    )
    parser.add_argument(
        "--protected",
        dest="protected",
        action="store_true",
        default=None,
        help="Protect image from deletion."
    )
    parser.add_argument(
        "--unprotected",
        dest="protected",
        action="store_false",
        default=None,
        help="Allow image to be deleted.",
    )
    parser.add_argument(
        "--public",
        dest="visibility",
        action="store_const",
        const="public",
        default=None,
        help="Make image accessible to all users."
    )
    parser.add_argument(
        "--private",
        dest="visibility",
        action="store_const",
        const="shared",
        default=None,
        help="Make image accessible only to the owners."
    )

class ListImage(base.Lister):
    _description = "List compute images."
    _default_fields = ['id', 'name', 'size', 'status', 'disk_format']
    _sort_keys = ['id', 'name', 'status', 'created_at', 'updated_at'
                  'size', 'disk_format']

    def configure_parser(self, parser):
        parser.add_argument(
            '--limit',
            metavar='<num>',
            type=int,
            help='The maximum number of images to list. To list all images, '
                 'set the option to -1.'
        )
        parser.add_argument(
            '--marker',
            metavar='<image>',
            help='List images after the marker.'
        )
        parser.add_argument(
            '--name',
            metavar='<name>',
            action='filter',
            operators='contains',
            help='List images with the specified name or use a filter.'
        )
        parser.add_argument(
            '--id',
            metavar='<id>',
            action='filter',
            operators='in',
            help='Show an image with the specified ID or list images using '
                 'a filter.'
        )
        parser.add_argument(
            '--status',
            metavar='<status>',
            action='filter',
            operators='in',
            help='List images with the specified status or use a filter.'
        )
        parser.add_argument(
            '--placement',
            metavar='<placement>',
            action='filter',
            operators='any',
            help='List images added to a placement with the specified ID or '
                 'use a filter'
        )
        parser.add_argument(
            '--disk-format',
            metavar='<disk-format>',
            help='List images with the specified disk format.'
        )
        parser.add_argument(
            '--project',
            metavar='<project>',
            action='filter',
            operators='in',
            help='List images that belong to projects with the specified names or IDs. '
                 'Can only be performed by system administrators. '
                 'Specify multiple project IDs as a comma-separated list.'
        )
        parser.add_argument(
            "--domain",
            metavar="<domain>",
            help='List images that belong to a domain with the specified name or ID. '
                 'Can only be performed by system administrators.'
        )
        parser.add_argument(
            '--sort',
            metavar='<sort>',
            help="List images sorted by key.\n"
                 "The sorting format is <sort-key>:<order>. The order is 'asc' or 'desc'.\n"
                 "Supported sort keys: {}".format(', '.join(self._sort_keys))
        )

    def do_action(self, parsed_args):
        filters = {}
        if parsed_args.name:
            filters['name'] = parsed_args.name
        if parsed_args.id:
            filters['id'] = parsed_args.id
        if parsed_args.status:
            filters['status'] = parsed_args.status
        if parsed_args.placement:
            filters['traits'] = parsed_args.placement
        if parsed_args.disk_format:
            filters['disk_format'] = parsed_args.disk_format
        if parsed_args.project:
            manager = self.app.vinfra.compute.projects
            filters['project_id'] = utils.validate_resources_from_operator(
                manager, parsed_args.project)
        if parsed_args.domain:
            domain = utils.find_resource(self.app.vinfra.domains, parsed_args.domain)
            filters['domain_id'] = domain.id
        if parsed_args.sort:
            filters['sort'] = parsed_args.sort

        data = self.app.vinfra.compute.images.list(
            limit=parsed_args.limit, marker=parsed_args.marker,
            filters=filters)
        return data

class ShowImage(base.ShowOne):
    _description = "Display compute image details."

    def configure_parser(self, parser):
        _image_arg(parser)

    def do_action(self, parsed_args):
        image = utils.find_resource(self.app.vinfra.compute.images,
                                    parsed_args.image)
        return image

class CreateImage(base.TaskCommand):
    _description = "Create a new compute image."

    def configure_parser(self, parser):
        _common_set_options(parser)

        disk_formats = ['aki', 'ami', 'ari', 'detect', 'iso', 'ploop',
                        'qcow2', 'raw', 'vdi', 'vhd', 'vhdx', 'vmdk']
        container_formats = [
            'aki', 'ami', 'ari', 'bare', 'docker', 'ovf', 'ova']

        parser.add_argument(
            "--disk-format",
            type=str,
            metavar="<disk_format>",
            default="detect",
            choices=disk_formats,
            help="Disk format {} (default: detect)".format(
                ','.join(disk_formats)
            )
        )
        parser.add_argument(
            "--container-format",
            type=str,
            metavar="<format>",
            default="bare",
            choices=container_formats,
            help="Container format: {} (default: bare)".format(
                ','.join(container_formats)
            )
        )
        parser.add_argument(
            "--tags",
            metavar="<tags>",
            type=parse_list_options,
            help="A comma-separated list of tags"
        )
        parser.add_argument(
            "name",
            metavar="<image-name>",
            help="Image name"
        )
        parser.add_argument(
            "--verify",
            action='store_true',
            help="Verify the checksum of the uploaded image"
        )
        parser.add_argument(
            "--file",
            metavar="<file>",
            required=True,
            help="Create image from a local file"
        )
        parser.add_argument(
            "--uefi",
            metavar="<uefi>",
            dest='hw_firmware_type',
            action='store_const',
            const='uefi',
            default=None,
            help="Create image with UEFI."
        )

    def do_action(self, parsed_args):
        def image_create(stream, args):
            return self.app.vinfra.compute.images.create_async(
                stream, args.name, args.disk_format, args.container_format,
                min_disk=args.min_disk, min_ram=args.min_ram,
                os_distro=args.os_distro, protected=args.protected,
                visibility=args.visibility, tags=args.tags,
                verify=args.verify,
                hw_firmware_type=args.hw_firmware_type,
            )

        if parsed_args.verify and not parsed_args.wait:
            raise exceptions.ValidationError(
                'The --verify option must be used together with the --wait option')

        if parsed_args.hw_firmware_type and parsed_args.file.endswith('.iso'):
            raise exceptions.ValidationError(
                "UEFI option cannot be used with ISO images."
            )

        # NOTE(akurbatov): request to upload image can get unauthorized
        # response. Client can then try to reauth in background and repeat
        # request. But stream offset must be then reset. Make sure user
        # passed correct login and password to not complicate logic by simple
        # backend API call:
        self.app.vinfra.get_meta()

        try:
            stream = open(parsed_args.file, mode='rb')
        except Exception as err:
            raise exceptions.ValidationError(
                'Failed to open {} ({})'.format(parsed_args.file, err))

        if (
                parsed_args.formatter != 'table'
                or not self.app.stderr.isatty()
        ):
            return image_create(stream, parsed_args)

        pattern = 'Uploading image to server [elapsed time: %s]... '
        widgets = [pb.Timer(format=pattern), pb.AnimatedMarker()]
        pbar = pb.ProgressBar(maxval=80, term_width=80, widgets=widgets,
                              fd=self.app.stderr)
        with utils.progress_bar_context(self.app, pbar):
            return image_create(stream, parsed_args)

class DeleteImage(base.Command):
    _description = "Delete a compute image."

    def configure_parser(self, parser):
        _image_arg(parser)

    def do_action(self, parsed_args):
        image = utils.find_resource(self.app.vinfra.compute.images,
                                    parsed_args.image)
        return image.delete()

class SetImage(base.ShowOne):
    _description = "Modify compute image parameters."

    def configure_parser(self, parser):
        _common_set_options(parser)
        parser.add_argument(
            "--name",
            metavar="<name>",
            help="Image name"
        )
        _image_arg(parser)

    def do_action(self, parsed_args):
        image = utils.find_resource(self.app.vinfra.compute.images,
                                    parsed_args.image)
        return image.update(
            name=parsed_args.name, min_disk=parsed_args.min_disk,
            min_ram=parsed_args.min_ram, os_distro=parsed_args.os_distro,
            protected=parsed_args.protected, visibility=parsed_args.visibility)

class SaveImage(base.Command):
    _description = "Download a compute image."

    def configure_parser(self, parser):
        parser.add_argument(
            "--file",
            metavar="<filename>",
            help="File to save the image to (default: stdout)"
        )
        _image_arg(parser)

    def do_action(self, parsed_args):
        image = utils.find_resource(self.app.vinfra.compute.images,
                                    parsed_args.image)
        fdst = None
        if not parsed_args.file:
            fdst = sys.stdout
        else:
            if os.path.exists(parsed_args.file):
                raise exceptions.ValidationError(
                    'File "{}" exists'.format(parsed_args.file))
            fdst = open(parsed_args.file, 'wb')

        image.download(fdst)


import logging
import multiprocessing
import warnings
from copy import copy
from typing import Dict, Optional, Tuple, Union, cast

import pandas as pd
import pyspark.sql.functions as F
from pandas import Series
from pyspark.ml import Transformer, PipelineModel
from pyspark.ml.util import MLWritable, MLReadable, MLWriter
from pyspark.ml.feature import VectorAssembler
from synapse.ml.lightgbm import LightGBMClassifier, LightGBMRegressor, LightGBMRegressionModel, LightGBMClassificationModel
from synapse.ml.onnx import ONNXModel
import lightgbm as lgb
from lightgbm import Booster, LGBMClassifier

from lightautoml.ml_algo.tuning.base import Distribution, SearchSpace
from lightautoml.pipelines.selection.base import ImportanceEstimator
from lightautoml.spark.dataset.base import SparkDataset
from lightautoml.spark.utils import SparkDataFrame
from lightautoml.spark.ml_algo.base import SparkTabularMLAlgo, SparkMLModel, AveragingTransformer
from lightautoml.spark.mlwriters import LightGBMModelWrapperMLReader, LightGBMModelWrapperMLWriter, ONNXModelWrapperMLReader, ONNXModelWrapperMLWriter
from lightautoml.spark.transformers.base import DropColumnsTransformer, PredictionColsTransformer, ProbabilityColsTransformer
from lightautoml.spark.validation.base import SparkBaseTrainValidIterator
from lightautoml.utils.timer import TaskTimer
from lightautoml.validation.base import TrainValidIterator

logger = logging.getLogger(__name__)

class LightGBMModelWrapper(Transformer, MLWritable, MLReadable):

    def __init__(self, model: Union[LightGBMRegressionModel, LightGBMClassificationModel] = None) -> None:
        super().__init__()
        self.model = model

    def write(self) -> MLWriter:
        return LightGBMModelWrapperMLWriter(self)

    @classmethod
    def read(cls):
        """Returns an MLReader instance for this class."""
        return LightGBMModelWrapperMLReader()

    def _transform(self, dataset: SparkDataset) -> SparkDataset:
        return self.model.transform(dataset)

class ONNXModelWrapper(Transformer, MLWritable, MLReadable):

    def __init__(self, model: ONNXModel = None) -> None:
        super().__init__()
        self.model = model

    def write(self) -> MLWriter:
        return ONNXModelWrapperMLWriter(self)

    @classmethod
    def read(cls):
        """Returns an MLReader instance for this class."""
        return ONNXModelWrapperMLReader()

    def _transform(self, dataset: SparkDataset) -> SparkDataset:
        return self.model.transform(dataset)

class SparkBoostLGBM(SparkTabularMLAlgo, ImportanceEstimator):

    _name: str = "LightGBM"

    _default_params = {
        "learningRate": 0.05,
        "numLeaves": 128,
        "featureFraction": 0.7,
        "baggingFraction": 0.7,
        "baggingFreq": 1,
        "maxDepth": -1,
        "minGainToSplit": 0.0,
        "maxBin": 255,
        "minDataInLeaf": 5,
        # e.g. num trees
        "numIterations": 3000,
        "earlyStoppingRound": 50,
        # for regression
        "alpha": 1.0,
        "lambdaL1": 0.0,
        "lambdaL2": 0.0,
    }

    # mapping between metric name defined via SparkTask
    # and metric names supported by LightGBM
    _metric2lgbm = {
        "binary": {
            "auc": "auc",
            "aupr": "areaUnderPR"
        },
        "reg": {
            "r2": "rmse",
            "mse": "mse",
            "mae": "mae",
        },
        "multiclass": {
            "crossentropy": "cross_entropy"
        }
    }

    def __init__(self,
                 cacher_key: str,
                 default_params: Optional[dict] = None,
                 freeze_defaults: bool = True,
                 timer: Optional[TaskTimer] = None,
                 optimization_search_space: Optional[dict] = {},
                 use_single_dataset_mode: bool = True,
                 max_validation_size: int = 10_000,
                 chunk_size: int = 4_000_000,
                 convert_to_onnx: bool = False,
                 mini_batch_size: int = 5000,
                 seed: int = 42):
        SparkTabularMLAlgo.__init__(self, cacher_key, default_params, freeze_defaults, timer, optimization_search_space)
        self._probability_col_name = "probability"
        self._prediction_col_name = "prediction"
        self._raw_prediction_col_name = "raw_prediction"
        self._assembler = None
        self._drop_cols_transformer = None
        self._use_single_dataset_mode = use_single_dataset_mode
        self._max_validation_size = max_validation_size
        self._seed = seed
        self._models_feature_impotances = []
        self._chunk_size = chunk_size
        self._convert_to_onnx = convert_to_onnx
        self._mini_batch_size = mini_batch_size

    def _infer_params(self) -> Tuple[dict, int]:
        """Infer all parameters in lightgbm format.

        Returns:
            Tuple (params, num_trees, early_stopping_rounds, verbose_eval, fobj, feval).
            About parameters: https://lightgbm.readthedocs.io/en/latest/_modules/lightgbm/engine.html

        """
        assert self.task is not None

        task = self.task.name

        params = copy(self.params)

        if "isUnbalance" in params:
            params["isUnbalance"] = True if params["isUnbalance"] == 1 else False

        verbose_eval = 1

        if task == "reg":
            params["objective"] = "regression"
            params["metric"] = self._metric2lgbm[task].get(self.task.metric_name, None)
        elif task == "binary":
            params["objective"] = "binary"
            params["metric"] = self._metric2lgbm[task].get(self.task.metric_name, None)
        elif task == "multiclass":
            params["objective"] = "multiclass"
            params["metric"] = "multiclass"
        else:
            raise ValueError(f"Unsupported task type: {task}")

        if task != "reg":
            if "alpha" in params:
                del params["alpha"]
            if "lambdaL1" in params:
                del params["lambdaL1"]
            if "lambdaL2" in params:
                del params["lambdaL2"]

        params = {**params}

        return params, verbose_eval

    def init_params_on_input(self, train_valid_iterator: TrainValidIterator) -> dict:
        self.task = train_valid_iterator.train.task

        sds = cast(SparkDataset, train_valid_iterator.train)
        rows_num = sds.data.count()
        task = train_valid_iterator.train.task.name

        suggested_params = copy(self.default_params)

        if self.freeze_defaults:
            # if user change defaults manually - keep it
            return suggested_params

        if task == "reg":
            suggested_params = {
                "learningRate": 0.05,
                "numLeaves": 32,
                "featureFraction": 0.9,
                "baggingFraction": 0.9,
            }

        if rows_num <= 10000:
            init_lr = 0.01
            ntrees = 3000
            es = 200

        elif rows_num <= 20000:
            init_lr = 0.02
            ntrees = 3000
            es = 200

        elif rows_num <= 100000:
            init_lr = 0.03
            ntrees = 2000
            es = 200
        elif rows_num <= 300000:
            init_lr = 0.04
            ntrees = 2000
            es = 100
        else:
            init_lr = 0.05
            ntrees = 2000
            es = 100

        if rows_num > 300000:
            suggested_params["numLeaves"] = 128 if task == "reg" else 244
        elif rows_num > 100000:
            suggested_params["numLeaves"] = 64 if task == "reg" else 128
        elif rows_num > 50000:
            suggested_params["numLeaves"] = 32 if task == "reg" else 64
            # params['reg_alpha'] = 1 if task == 'reg' else 0.5
        elif rows_num > 20000:
            suggested_params["numLeaves"] = 32 if task == "reg" else 32
            suggested_params["alpha"] = 0.5 if task == "reg" else 0.0
        elif rows_num > 10000:
            suggested_params["numLeaves"] = 32 if task == "reg" else 64
            suggested_params["alpha"] = 0.5 if task == "reg" else 0.2
        elif rows_num > 5000:
            suggested_params["numLeaves"] = 24 if task == "reg" else 32
            suggested_params["alpha"] = 0.5 if task == "reg" else 0.5
        else:
            suggested_params["numLeaves"] = 16 if task == "reg" else 16
            suggested_params["alpha"] = 1 if task == "reg" else 1

        suggested_params["learningRate"] = init_lr
        suggested_params["numIterations"] = ntrees
        suggested_params["earlyStoppingRound"] = es

        if task != "reg":
            if "alpha" in suggested_params:
                del suggested_params["alpha"]

        return suggested_params

    def _get_default_search_spaces(self, suggested_params: Dict, estimated_n_trials: int) -> Dict:
        """Train on train dataset and predict on holdout dataset.

        Args:
            fold_prediction_column: column name for predictions made for this fold
            full: Full dataset that include train and valid parts and a bool column that delimits records
            train: Train Dataset.
            valid: Validation Dataset.

        Returns:
            Target predictions for valid dataset.

        """
        assert self.task is not None

        optimization_search_space = dict()

        optimization_search_space["featureFraction"] = SearchSpace(
            Distribution.UNIFORM,
            low=0.5,
            high=1.0,
        )

        optimization_search_space["numLeaves"] = SearchSpace(
            Distribution.INTUNIFORM,
            low=4,
            high=255,
        )

        if self.task.name == "binary" or self.task.name == "multiclass":
            optimization_search_space["isUnbalance"] = SearchSpace(
                Distribution.DISCRETEUNIFORM,
                low=0,
                high=1,
                q=1
            )

        if estimated_n_trials > 30:
            optimization_search_space["baggingFraction"] = SearchSpace(
                Distribution.UNIFORM,
                low=0.5,
                high=1.0,
            )

            optimization_search_space["minSumHessianInLeaf"] = SearchSpace(
                Distribution.LOGUNIFORM,
                low=1e-3,
                high=10.0,
            )

        if estimated_n_trials > 100:
            if self.task.name == "reg":
                optimization_search_space["alpha"] = SearchSpace(
                    Distribution.LOGUNIFORM,
                    low=1e-8,
                    high=10.0,
                )

            optimization_search_space["lambdaL1"] = SearchSpace(
                Distribution.LOGUNIFORM,
                low=1e-8,
                high=10.0,
            )

        return optimization_search_space

    def predict_single_fold(self,
                            dataset: SparkDataset,
                            model: Union[LightGBMRegressor, LightGBMClassifier]) -> SparkDataFrame:

        temp_sdf = self._assembler.transform(dataset.data)

        pred = model.transform(temp_sdf)

        return pred

    def fit_predict_single_fold(self,
                                fold_prediction_column: str,
                                full: SparkDataset,
                                train: SparkDataset,
                                valid: SparkDataset) -> Tuple[SparkMLModel, SparkDataFrame, str]:
        assert self.validation_column in full.data.columns, 'Train should contain validation column'

        if self.task is None:
            self.task = full.task

        (
            params,
            verbose_eval
        ) = self._infer_params()

        logger.info(f"Input cols for the vector assembler: {full.features}")
        logger.info(f"Running lgb with the following params: {params}")

        if self._assembler is None:
            self._assembler = VectorAssembler(
                inputCols=self.input_features,
                outputCol=f"{self._name}_vassembler_features",
                handleInvalid="keep"
            )

        LGBMBooster = LightGBMRegressor if full.task.name == "reg" else LightGBMClassifier

        if full.task.name in ['binary', 'multiclass']:
            params['rawPredictionCol'] = self._raw_prediction_col_name
            params['probabilityCol'] = fold_prediction_column
            params['predictionCol'] = self._prediction_col_name
            params['isUnbalance'] = True
        else:
            params['predictionCol'] = fold_prediction_column

        master_addr = train.spark_session.conf.get('spark.master')
        if master_addr.startswith('local'):
            cores_str = master_addr[len("local["):-1]
            cores = int(cores_str) if cores_str != "*" else multiprocessing.cpu_count()
            params["numThreads"] = max(cores - 1, 1)
        else:
            params["numThreads"] = max(int(train.spark_session.conf.get("spark.executor.cores", "1")) - 1, 1)

        train_data = full.data
        valid_size = train_data.where(F.col(self.validation_column) == 1).count()
        max_val_size = self._max_validation_size
        if valid_size > max_val_size:
            warnings.warn(f"Maximum validation size for SparkBoostLGBM is exceeded: {valid_size} > {max_val_size}. "
                          f"Reducing validation size down to maximum.", category=RuntimeWarning)
            rest_cols = list(train_data.columns)
            rest_cols.remove(self.validation_column)

            replace_col = F.when(F.rand(self._seed) < max_val_size / valid_size, F.lit(True)).otherwise(F.lit(False))
            val_filter_cond = F.when(F.col(self.validation_column) == 1, replace_col).otherwise(F.lit(True))

            train_data = train_data.where(val_filter_cond)

        valid_data = valid.data

        lgbm = LGBMBooster(
            **params,
            featuresCol=self._assembler.getOutputCol(),
            labelCol=full.target_column,
            validationIndicatorCol=self.validation_column,
            verbosity=verbose_eval,
            useSingleDatasetMode=self._use_single_dataset_mode,
            isProvideTrainingMetric=True,
            chunkSize=self._chunk_size
        )

        logger.info(f"Use single dataset mode: {lgbm.getUseSingleDatasetMode()}. NumThreads: {lgbm.getNumThreads()}")

        if full.task.name == "reg":
            lgbm.setAlpha(0.5).setLambdaL1(0.0).setLambdaL2(0.0)

        ml_model = lgbm.fit(self._assembler.transform(train_data))

        val_pred = ml_model.transform(self._assembler.transform(valid_data))
        val_pred = DropColumnsTransformer(
            remove_cols=[],
            optional_remove_cols=[self._prediction_col_name, self._probability_col_name, self._raw_prediction_col_name]
        ).transform(val_pred)

        self._models_feature_impotances.append(ml_model.getFeatureImportances(importance_type='gain'))

        if self._convert_to_onnx:
            logger.info("Model convert is started")
            booster_model_str = ml_model.getLightGBMBooster().modelStr().get()
            booster = lgb.Booster(model_str=booster_model_str)
            model_payload_ml = self._convertModel(booster, len(self.input_features))

            onnx_ml = ONNXModel().setModelPayload(model_payload_ml)

            if full.task.name == "reg":
                onnx_ml = (
                    onnx_ml
                        .setDeviceType("CPU")
                        .setFeedDict({"input": f"{self._name}_vassembler_features"})
                        .setFetchDict({ml_model.getPredictionCol(): "variable"})
                        .setMiniBatchSize(self._mini_batch_size)
                )
            else:
                onnx_ml = (
                    onnx_ml
                        .setDeviceType("CPU")
                        .setFeedDict({"input": f"{self._name}_vassembler_features"})
                        .setFetchDict({ml_model.getProbabilityCol(): "probabilities", ml_model.getPredictionCol(): "label"})
                        .setMiniBatchSize(self._mini_batch_size)
                )

            ml_model = onnx_ml
            logger.info("Model convert is ended")

        return ml_model, val_pred, fold_prediction_column

    def fit(self, train_valid: SparkBaseTrainValidIterator):
        logger.info("Starting LGBM fit")
        self.fit_predict(train_valid)
        logger.info("Finished LGBM fit")

    def get_features_score(self) -> Series:
        imp = 0
        for model_feature_impotances in self._models_feature_impotances:
            imp = imp + pd.Series(model_feature_impotances)

        imp = imp / len(self._models_feature_impotances)

        result = Series(list(imp), index=self.features).sort_values(ascending=False)
        return result

    @staticmethod
    def _convertModel(lgbm_model: Booster, input_size: int) -> bytes:
        from onnxmltools.convert import convert_lightgbm
        from onnxconverter_common.data_types import FloatTensorType
        initial_types = [("input", FloatTensorType([-1, input_size]))]
        onnx_model = convert_lightgbm(lgbm_model, initial_types=initial_types, target_opset=9)
        return onnx_model.SerializeToString()

    def _build_transformer(self) -> Transformer:
        avr = self._build_averaging_transformer()
        if self._convert_to_onnx:
            wrapped_models = [ONNXModelWrapper(m) for m in self.models]
        else:
            wrapped_models = [LightGBMModelWrapper(m) for m in self.models]
        models = [el for m in wrapped_models for el in [m, DropColumnsTransformer(
                remove_cols=[],
                optional_remove_cols=[self._prediction_col_name,
                                    self._probability_col_name,
                                    self._raw_prediction_col_name]
            )
        ]]
        if self._convert_to_onnx:
            if self.task.name in ['binary', 'multiclass']:
                models.append(ProbabilityColsTransformer(probability_сols=self._models_prediction_columns, num_classes=self.n_classes))
            else:
                models.append(PredictionColsTransformer(prediction_сols=self._models_prediction_columns))
        averaging_model = PipelineModel(stages=[self._assembler] + models + [avr])
        return averaging_model

    def _build_averaging_transformer(self) -> Transformer:
        avr = AveragingTransformer(
            self.task.name,
            input_cols=self._models_prediction_columns,
            output_col=self.prediction_feature,
            remove_cols=[self._assembler.getOutputCol()] + self._models_prediction_columns,
            convert_to_array_first=not (self.task.name == "reg"),
            dim_num=self.n_classes
        )
        return avr

    def fit_predict(self, train_valid_iterator: SparkBaseTrainValidIterator) -> SparkDataset:
        """Fit and then predict accordig the strategy that uses train_valid_iterator.

        If item uses more then one time it will
        predict mean value of predictions.
        If the element is not used in training then
        the prediction will be ``numpy.nan`` for this item

        Args:
            train_valid_iterator: Classic cv-iterator.

        Returns:
            Dataset with predicted values.

        """
        logger.info("Starting LGBM fit")
        self.timer.start()

        self.input_roles = train_valid_iterator.input_roles

        res = super().fit_predict(train_valid_iterator)

        logger.info("Finished LGBM fit")
        return res


import re, psycopg2, os
from fastai.text import load_learner
from pathlib import Path
from dotenv import load_dotenv
load_dotenv()

path = Path(__file__).parent
learn = load_learner("./", path/'models/deep_poet')
URI = os.getenv("URI")
def generate_poem(init,length,temp):
  return learn.predict(init, n_words=length, temperature=temp)

def getMax(poem):
  poem = re.sub('\n',"",poem)
  list = poem.split(" ")
  occ = {}
  max = 0
  for word in list:
    if len(word) > 3  :
      occ[word] = occ.get(word,0) + 1
      if occ[word] > max : max = occ[word]
  return max

def savePoem(poem, init, length, temp, counter, repetition_tolerance):
  connection = psycopg2.connect(URI) #.env uri
  cursor = connection.cursor()
  args = [poem, init, length, temp, counter, repetition_tolerance]
  cursor.execute("INSERT INTO poems VALUES (default, %s, %s, %s, %s, %s, %s)", args)
  connection.commit()

def write_poem(init, temp=0.55, length=100, repetition_tolerance=8, max=10):
  counter = 0
  poem = generate_poem(init,length,temp)
  while getMax(poem)>= repetition_tolerance :
      counter += 1
      poem = generate_poem(init,length,temp)
      if (counter > max) : break
  poem = re.sub(' i ', ' I ', poem)
  verses = poem.split("\n")
  for i in range(len(verses)):
      verses[i] = verses[i].strip()
      if len(verses[i]) > 1: verses[i] = verses[i][0].upper() + verses[i][1:]
      elif len(verses[i]) == 1 : verses[i] = verses[i].upper()
  poem = "<br>".join(verses)
  savePoem(poem, init, length, temp, counter, repetition_tolerance)
  return poem


from PIL import Image

def convert_image_to_hex(image_path, output_path):
    # Load the image file
    img = Image.open(image_path)

    # Resize the image to 16x16 if it's not already that size
    if img.size != (16, 16):
        img = img.resize((16, 16), Image.ANTIALIAS)

    # Convert the image to an RGB array
    rgb_img = img.convert('RGB')

    # Function to convert 8-bit color depth to 4-bit color depth
    def to_4bit(value):
        return value // 16

    # Create an empty list to hold the 12-bit color values for the image
    hex_values = []

    # Loop over each pixel and convert to 12-bit color depth
    for y in range(16):
        for x in range(16):
            r, g, b = rgb_img.getpixel((x, y))
            r, g, b = to_4bit(r), to_4bit(g), to_4bit(b)
            # Combine the 4-bit R, G, and B values into a single 12-bit number
            color_12bit = (r << 8) | (g << 4) | b
            # Format the number as a hex string and append to our list
            hex_values.append(format(color_12bit, '03x'))

    # Save the list of color values to a hex file
    with open(output_path, 'w') as file:
        for hex_value in hex_values:
            file.write(f"{hex_value}\n")

convert_image_to_hex('/Users/aamaya3/Desktop/image_converter/crouch.png', '/Users/aamaya3/Desktop/image_converter/crouch.hex')


# 공간 복합도 : O(n)
# 시간 복잡도 : O(n log n)

#helper function이라고도 부름
# [1,3,7,8] [2,4,5,6]

def merge(list1, list2):
    combined = []
    i = 0
    j = 0
    while i < len(list1) and j < len(list2):
        if list1[i] < list2[j]:
            combined.append(list1[i])
            i += 1
        else:
            combined.append(list2[j])
            j += 1

    while i < len(list1):
        combined.append(list1[i])
        i += 1

    while j < len(list2):
        combined.append(list2[j])
        j += 1

    return combined

#print(merge([1,3,7,8], [2,4,5,6]))

"""
[1,2,3,4,5,6,7,8]

1)Breaks lists in half [1,2]
2)Base case: when len(the_list) is 1
3)Uses merge() to put lists together
"""

def merge_sort(my_list):
    if len(my_list) == 1:
        return my_list
    # 리스트 안의 값이 홀수의 경우 나누면 소수가 발생
    mid = int(len(my_list)/2)
    left = my_list[:mid]
    right = my_list[mid:]
    #이렇게만 할경우 각 리스트가 정렬되지 않음 -> merge적용 어려움
    #return merge(left, right)
    return merge(merge_sort(left), merge_sort(right))

print(merge_sort([3,1,4,2]))


from pprint import pprint

from googleapiclient import discovery

import google.auth
import re
import sys
import os

credentials, project = google.auth.default()
service = discovery.build('compute', 'v1', credentials=credentials)

#print(str(sys.argv))
if (len(sys.argv)) >=3:
    project = sys.argv[1]
    instance = sys.argv[2]
    zone = sys.argv[3]

    # Get instance details
    request = service.instances().get(project=project, zone=zone, instance=instance)
    response = request.execute()

    # Get machine type
    mtype = re.search(r'(.*)/(.*)', response['machineType']).group(2)

    # Get operating system name
    osu = response['disks'][0]['licenses'][0]
    os = re.search(r'(.*)/(.*)', osu).group(2)

    # Get disk size/count
    #dsize = str(response['disks'][0]['diskSizeGb']) + ' GB'
    num_disks = len(response['disks'])

    # Use machine type to get cpu count & memory size
    mrequest = service.machineTypes().get(project=project, zone=zone, machineType=mtype)
    mresponse = mrequest.execute()

    # Get cpu count
    cpu = mresponse['guestCpus']

    # Get memory size
    megabyte = mresponse['memoryMb']
    gigabyte = 1.0/1024
    memory = str(gigabyte * megabyte) + ' GB'

    # Get status
    status = response['status']

    # Check if the instance is running
    if status == 'RUNNING':
        publicip_status = "IP externo desativado"
        publicip = ""
    else:
        publicip_status = "VM desligada"
        publicip = ""

    # Get network Interfaces list of dictionaries
    ni = response['networkInterfaces']

    # Get private IP and public IP
    for d in ni:
        privateip = str(d['networkIP'])
        if 'accessConfigs' in d:
            if 'natIP' in d['accessConfigs'][0]:
                publicip = str(d['accessConfigs'][0]['natIP'])
                publicip_status = "IP reservado"
            else:
                publicip = "IP temporario"
        else:
            publicip = "NA"

    print(f'{project},{instance},{zone},{mtype},{status},{os},{cpu},{memory},{num_disks},{privateip},{publicip_status},{publicip}')


## 광주인공지능사관학교 Warm-up Project (Team - DOYA)
## 커리어넷에서 전공 관련된 정보들을 크롤링하자!
## 2020.09.09

import requests
from bs4 import BeautifulSoup, NavigableString, Tag
import csv

major_link = [] # 각 전공 별로의 href를 할당하기 위한 list
major_big = []

for page in range(1, 53):
    URL = f'https://www.career.go.kr/cnet/front/base/major/FunivMajorList.do?pageIndex={page}' # 1~52 페이지까지 있는 커리어넷의 홈페이지 특성에 따라 URL 할당

    response = requests.get(URL)
    soup = BeautifulSoup(response.text, 'html.parser')

    soup_section = soup.select('#view1 > ul > li')

    for data in soup_section:
        a_tag = data.select_one('a')['href']
        major_link.append(a_tag) # 1~52페이지까지 돌면서 각 전공들의 구체적인 내용들이 담겨 있는 href를 가지고 와서 major_link라는 리스트에 저장
        # len(major_link)는 516개 = 커리어넷에 있는 전공의 개수

major_dic = { # csv 파일에 넣기 위해서 dic를 만듦
    'major': "",
    'major_describe': ""
}

for link in major_link: # 각 전공의 link에 들어가서 데이터들을 가져옴
    major_small = [] # 초기화

    URL = 'https://www.career.go.kr' + link

    response = requests.get(URL)
    soup = BeautifulSoup(response.text, 'html.parser')

    title = soup.select_one('#contents > #frm > .word_tit > h2').text
    detail = soup.select('#tab1 > .word_ul')

    for i in detail:
        major_small.append(i.text.replace('\n', '').replace('\r', '')) # 전공과 관련된 설명 부분들을 모두 모아서 major_small 이라는 작은 list에 저장

    major_dic['major'] = title # 전공 이름
    major_dic['major_describe'] = major_small # 전공 설명

    with open('careernet_major.csv', 'a', newline='') as csvfile: # for문이 한번 실행 될때마다 하나의 전공을 한 행씩 저장하는 csv 파일 생성
        fieldnames = ['major', 'major_describe']
        csv_writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        csv_writer.writerow(major_dic)

    print(111)


"""Class for sending metrics to ElasticSearch ci-performance-test index."""
import datetime
import json
import time

import gflags
import requests

gflags.DEFINE_boolean("send_perf_metrics", False, "Send perf metrics to ElasticSearch for dashboard presentation.")

FLAGS = gflags.FLAGS
BASE_URL = "http://elasticsearch.dfinity.systems:9200"
INDEX = "performance-trend"

class ElasticSearch:
    """Client class for intercting with ElasticSearch."""

    def call_elastic_search(metric):
        """Client to make a call to ElasticSearch."""
        addr = f"{BASE_URL}/{INDEX}/_doc/"
        for i in range(3):
            try:
                response = requests.post(
                    url=addr,
                    data=json.dumps(metric, indent=4, default=str).encode(),
                    headers={"content-type": "application/json"},
                    timeout=30,
                )
                if response.status_code < 500:
                    if response.status_code < 400:
                        # success
                        print(f"Data points successfully sent to {addr}")
                        return True
                    else:
                        # client side error, log error and no retry
                        print(
                            f"Invalid input request: {metric}. Response received was: {response.status_code} {response.content}"
                        )
                        return False

                else:
                    # service side error, retry after 3 seconds
                    print(
                        f'Remote service error received when send metrics to {addr}. {response.status_code} {response.content}. {"" if (i == 2) else "Sleep 3 seconds before next retry."}'
                    )
                    time.sleep(3)

            except Exception as e:
                print(
                    f'Exception caught when send metrics to {addr}. {e}. {"" if (i == 2) else "Sleep 3 seconds before next retry."}'
                )
                time.sleep(3)
        return False

    def send_perf(name, is_success, request_type, revision, branch, is_ci_job, summaries):
        """Send a performance nightly data point performance-trend index in ElasticSearch."""
        if FLAGS.send_perf_metrics is False:
            return True

        (
            failure_rate,
            failure_rate_threshold,
            median_latency,
            median_latency_threshold,
            throughput,
            throughput_threshold,
        ) = summaries

        metric = {
            "timestamp": datetime.datetime.now().isoformat(),
            "rev": revision,
            "branch": branch,
            "is_ci_job": is_ci_job,
            "package": "replica-perf-trend",
            "performance": {
                "title": name,
                "request_type": request_type,
                "is_success": is_success,
                "metrics": {
                    "failure_rate": {
                        "value": failure_rate,
                        "threshold": failure_rate_threshold,
                    },
                    "median_latency": {
                        "value": median_latency,
                        "threshold": median_latency_threshold,
                    },
                    "throughput": {
                        "value": throughput,
                        "threshold": throughput_threshold,
                    },
                },
            },
        }
        return ElasticSearch.call_elastic_search(metric)

    def send_max_capacity(
        name, request_type, experiment_revision, experiment_branch, is_ci_job, rps, out_dir, unix_time=None
    ):
        """Send a maximum capacity data point to performance-trend index in ElasticSearch."""
        if FLAGS.send_perf_metrics is False:
            return True

        if unix_time is None:
            unix_time = int(time.time())
            timestamp = datetime.datetime.fromtimestamp(unix_time).isoformat()
        else:
            timestamp = datetime.datetime.fromtimestamp(unix_time).isoformat()

        metric = {
            "unix_time": unix_time,
            "timestamp": timestamp,
            "rev": experiment_revision,
            "branch": experiment_branch,
            "is_ci_job": is_ci_job,
            "package": "replica-perf-trend",
            "report_url": f"http://10.11.10.88:9099/{out_dir}/report.html",
            "performance": {"title": name, "request_type": request_type, "metrics": {"max-capacity": rps}},
        }

        return ElasticSearch.call_elastic_search(metric)


#Magic Function
class person:
    nationality="Indian"#Class level encapsulation
    def sayhi(self):
        print("Hi "+self.fname+" "+self.lname)
    def __init__(obj,f,l):#initialiser,constructor
        obj.fname=f
        obj.lname=l
p1=person("Sachin","Tendulkar")
p2=person("Rahul","Dravid")
p1.sayhi()
p2.sayhi()


class Solution:
    def cherryPickup(self, grid: List[List[int]]) -> int:
        n = len(grid)
        @cache
        def dp(r1,c1,r2,c2):
            if r1 >= n or r2 >= n or c1 >= n or c2 >= n or grid[r1][c1] == -1 or grid[r2][c2] == -1 :
                return -float('inf')
            if r1 == n-1 and c1 == n- 1:
                return grid[r1][c1]
            if r2 == n-1 and c2 == n-1:
                return grid[r2][c2]
            cherries = 0
            if r1 == r2 and c1 == c2 :
                cherries = grid[r1][c1]
            else:
                cherries = grid[r1][c1] + grid[r2][c2]

            cherries += max(dp(r1+1,c1,r2+1,c2),dp(r1+1,c1,r2,c2+1),dp(r1,c1+1,r2+1,c2),dp(r1,c1+1,r2,c2+1))
            return cherries
        return max(0,dp(0,0,0,0))


import numpy as np
import matplotlib.pyplot as plt

#plt.show(block=True) close figure to continue code.
N=np.array([100,200,400])

for n in N:

        K = np.linspace(0,6,50)
        t = np.linspace(0,10,1000)

        x1 = np.zeros([n,len(t)])

        R = np.zeros(len(K))

        #omegas for n oscillators
        w = np.random.standard_cauchy(n)

        #initial condition for oscillators
        x1[:,0] = np.random.uniform(0,2*np.pi,n)
        h = t[3]-t[2]

        def r(x):
            r1 = np.absolute(np.sum(np.exp(x*1j)))/n
            w1 = np.angle(np.sum(np.exp(x*1j)))/n
            return r1,w1

        def f1(y):
                a,b=r(y)
                return w + (k*a)*np.sin(b-y)

        for j in range(len(K)):
            k = K[j]
            for i in range(len(t)-1):
                x = x1[:,i]
                k1 = h*f1(x)
                k2 = h*f1(x+k1/2.0)
                k3 = h*f1(x+k2/2.0)
                k4 = h*f1(x+k3)
                t[i+1] = t[i] + h
                x1[:,i+1] = x + ((1.0/6)*(k1+2*k2+2*k3+k4))

            a1,b1=r(x1[:,len(t)-30])
            R[j]=a1

        plt.scatter(K,R)
        plt.title('r v/s K for N='+str(n))
        plt.xlabel('K')
        plt.ylabel('|r|')
        plt.show()


import tkinter as tk
from tkinter import ttk
win=tk.Tk()
win.title('ADD GUI')
'''
fm=Frame(win)
fm.pack(fill=BOTH)
Button(fm,text="Button1",fg="red").pack()
Button(fm,text="Button2",fg="blue").pack()
Button(fm,text="Button3",fg="yellow").pack(side=LEFT)
'''
# create labels
name_lable=ttk.Label(win,text="Enter your name:")
name_lable.grid(row=0,column=0,sticky=tk.W)

contact_lable=ttk.Label(win,text="Enter your contact:")
contact_lable.grid(row=1,column=0,sticky=tk.W)

Email_lable=ttk.Label(win,text="Enter your Email:")
Email_lable.grid(row=3,column=0,sticky=tk.W)

Dept_lable=ttk.Label(win,text="Enter your Dept:")
Dept_lable.grid(row=4,column=0,sticky=tk.W)

# create entry box
name_var = tk.string()
name_entrybox=ttk.Entry(win,width=16,textvariable=name_var)
name_entrybox.grid(row=0,column=1)
'''
contact_var=tk.stringVar()
contact_entrybox=ttk.Entry(win,width=16,textvariable=contact_var)
contact_entrybox.grid(row=1,column=1)

Email_var=tk.stringVar()
Email_entrybox=ttk.Entry(win,width=16,textvariable=Email_var)
Email_entrybox.grid(row=2,column=1)

dep_var=tk.stringVar()
dept_entrybox=ttk.Entry(win,width=16,textvariable=dept_var)
dept_entrybox.grid(row=3,column=1)
'''

win.mainloop()


from pyftpdlib.filesystems import AbstractedFS, FilesystemError
from io import BytesIO
import django
import sys
import os

# 将项目路径添加到系统搜寻路径当中，查找方式为从当前脚本开始，找到要调用的django项目的路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
# 设置项目的配置文件 不做修改的话就是 settings 文件
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "webserver.settings")
django.setup()  # 加载项目配置

from api.harbor import FtpHarborManager
from api.exceptions import HarborError

class HarborFileSystem(AbstractedFS):
    def __init__(self, *args, **kwargs):
        super(HarborFileSystem, self).__init__(*args, **kwargs)
        self.bucket_name = self.root
        self.root = '/'
        self.client = FtpHarborManager()

    def realpath(self, path):
        return path

    def isdir(self, path):
        try:
            return self.client.ftp_is_dir(self.bucket_name, path.lstrip('/'))
        except HarborError as error:
            return False

    def isfile(self, path):
        try:
            return self.client.ftp_is_file(self.bucket_name, path.lstrip('/'))
        except HarborError as error:
            return False

    def islink(self, fs_path):
        return False

    def chdir(self, path):
        if self.isdir(path):
            self._cwd = self.fs2ftp(path)
        else:
            raise FilesystemError('Not a dir.')

    def stat(self, path):
        pass

    def lstat(self, path):
        pass

    def listdir(self, path):
        dir_list = []
        try:
            files = []
            files_generator = self.client.ftp_list_dir_generator(self.bucket_name, path[1:], per_num=2000)
            for _ in range(10):
                try:
                    data = next(files_generator)
                except StopIteration as error:
                    break
                files.append(data)
            for fi in files:
                for file in fi:
                    if file.fod is True:
                        if file.upt:
                            dir_list.append((file.name, file.upt, file.si))
                        else:
                            dir_list.append((file.name, file.ult, file.si))
                    else:
                        dir_list.append((file.name + '/', file.ult, 0))
        except HarborError as error:
            raise FilesystemError(error.msg)
        return dir_list

    def format_list(self, basedir, listing, ignore_err=True):
        assert isinstance(basedir, str), basedir

        # path = self.fs2ftp(basedir)
        for tmp in listing:
            if isinstance(tmp, tuple):
                filename, mtimestr, size = tmp
                if filename.endswith('/'):
                    perm = "drwxrwxrwx"
                    mtimestr = mtimestr.strftime('%b %d %X')[:-3]
                    filename = filename[:-1]
                else:
                    mtimestr = mtimestr.strftime('%b %d %X')[:-3]
                    perm = "-r-xr-xr-x"
                line = "%s %3s %-8s %-8s %8s %12s %s\r\n" % (
                    perm, 1, 'root', 'root', size, mtimestr, filename)

                if self.cmd_channel is not None:
                    yield line.encode("utf8", self.cmd_channel.unicode_errors)
                else:
                    yield line.encode("utf8")
            else:
                perm = "-r-xr-xr-x"
                line = "%s %3s %-8s %-8s %8s %12s %s\r\n" % (
                    perm, 1, 'root', 'root', '', '', tmp)

                if self.cmd_channel is not None:
                    yield line.encode("utf8", self.cmd_channel.unicode_errors)
                else:
                    yield line.encode("utf8")

    def format_mlsx(self, basedir, listing, perms, facts, ignore_err=True):
        assert isinstance(basedir, str), basedir

        ftp_path = self.fs2ftp(basedir)
        if len(listing) == 1:
            if isinstance(listing[0], tuple):
                if listing[0][0].endswith('/'):
                    _type = "dir"
                    perm = 'r'
                    filename = listing[0][0][:-1]
                    mtimestr = listing[0][1]
                    mtimestr = str(mtimestr).split('.')[0].replace('-', '').replace(':', '').replace(' ', '')
                    size = listing[0][2]
                    line = "type=%s;size=%d;perm=%s;modify=%s;unique=%s; %s\r\n" % (
                        _type, size, perm, mtimestr, '', filename)
                    if self.cmd_channel is not None:
                        yield line.encode("utf8", self.cmd_channel.unicode_errors)
                    else:
                        yield line.encode("utf8")
                else:
                    filename = listing[0][0]
                    mtimestr = listing[0][1]
                    mtimestr = str(mtimestr).split('.')[0].replace('-', '').replace(':', '').replace(' ', '')
                    size = listing[0][2]
                    perm = 'el'
                    _type = "file"
                    line = "type=%s;size=%d;perm=%s;modify=%s;unique=%s; %s\r\n" % (
                        _type, size, perm, mtimestr, '', filename)
                    if self.cmd_channel is not None:
                        yield line.encode("utf8", self.cmd_channel.unicode_errors)
                    else:
                        yield line.encode("utf8")
            else:
                ftp_path = os.path.join(ftp_path, listing[0])
                try:
                    data = self.client.ftp_get_obj(self.bucket_name, ftp_path[1:])
                    filename = data.name
                    mtimestr = data.upt
                    mtimestr = str(mtimestr).split('.')[0].replace('-', '').replace(':', '').replace(' ', '')
                    size = data.si
                    perm = 'el'
                    _type = "file"
                    line = "type=%s;size=%d;perm=%s;modify=%s;unique=%s; %s\r\n" % (
                        _type, size, perm, mtimestr, '', filename)
                    if self.cmd_channel is not None:
                        yield line.encode("utf8", self.cmd_channel.unicode_errors)
                    else:
                        yield line.encode("utf8")
                except HarborError as error:
                    raise FilesystemError(error.msg)
        else:
            for filename, mtimestr, size in listing:
                if filename.endswith('/'):
                    _type = "dir"
                    perm = 'r'
                    filename = filename[:-1]
                else:
                    perm = 'el'
                    _type = "file"
                mtimestr = str(mtimestr).split('.')[0].replace('-', '').replace(':', '').replace(' ', '')
                line = "type=%s;size=%d;perm=%s;modify=%s;unique=%s; %s\r\n" % (
                    _type, size, perm, mtimestr, '', filename)

                if self.cmd_channel is not None:
                    yield line.encode("utf8", self.cmd_channel.unicode_errors)
                else:
                    yield line.encode("utf8")

    def open(self, filename, mode):
        """Open a file returning its handler."""
        assert isinstance(filename, str), filename
        # print('function: open', filename, mode)
        ftp_path = self.fs2ftp(filename)
        mode = mode.lower()
        return FileHandler(self.bucket_name, ftp_path, self.client, mode)

    def mkdir(self, path):
        ftp_path = self.fs2ftp(path)
        try:
            self.client.ftp_mkdir(self.bucket_name, ftp_path[1:])
        except (HarborError, Exception) as error:
            raise FilesystemError(str(error))

    def rename(self, src, dst):
        new_name = os.path.basename(dst)
        new_dir = os.path.dirname(dst)
        try:
            self.client.ftp_move_rename(self.bucket_name, src[1:], new_name, new_dir)
        except HarborError as error:
            raise FilesystemError(f'rename dir is not supported, {str(error)}')
        except Exception as error:
            raise FilesystemError(str(error))

    def lexists(self, path):
        ftp_path = self.fs2ftp(path)

        if self.isdir(ftp_path) or self.isfile(ftp_path):
            return True
        else:
            return False

    def rmdir(self, path):
        ftp_path = self.fs2ftp(path)
        try:
            self.client.ftp_rmdir(self.bucket_name, ftp_path[1:])
        except (HarborError, Exception) as error:
            raise FilesystemError(str(error))

    def remove(self, path):
        ftp_path = self.fs2ftp(path)
        try:
            self.client.ftp_delete_object(self.bucket_name, ftp_path[1:])
        except (HarborError, Exception) as error:
            raise FilesystemError(str(error.msg))

    def getsize(self, path):
        ftp_path = self.fs2ftp(path)
        try:
            return self.client.ftp_get_obj_size(self.bucket_name, ftp_path[1:])
        except (HarborError, Exception) as error:
            raise FilesystemError(str(error.msg))

class FileHandler(object):
    def __init__(self, bucket_name, ftp_path, client, mode):
        self.bucket_name = bucket_name
        self.name = os.path.basename(ftp_path)
        self.ftp_path = ftp_path
        self.client = client
        self.closed = False
        self.file = BytesIO()
        self.offset = 0             # file pointer position
        self.is_breakpoint = False  # 标记是否断点续传
        self.write_generator = None
        self.read_generator = None
        self.mode = mode

    def ensure_init_write_generator(self, is_break_point=None):
        """
        确保已初始化 写生成器
        """
        if self.write_generator:
            return

        is_break_point = self.is_breakpoint if is_break_point is None else is_break_point
        try:
            self.write_generator = self.client.ftp_get_write_generator(
                self.bucket_name, self.ftp_path[1:], is_break_point)
            next(self.write_generator)
        except HarborError as error:
            raise FilesystemError(error.msg)

    def ensure_init_read_generator(self):
        if self.read_generator:
            return

        try:
            self.read_generator, ob = self.client.ftp_get_obj_generator(
                self.bucket_name, self.ftp_path[1:], offset=self.offset, per_size=4 * 1024 ** 2)
        except HarborError as error:
            raise FilesystemError(error.msg)

    def write(self, data):
        self.ensure_init_write_generator()
        self.file.write(data)
        if self.file.tell() >= 1024 ** 2 * 64:
            self._sync_cache()

        return len(data)

    def read(self, size=None):
        self.ensure_init_read_generator()
        try:
            data = next(self.read_generator)
        except Exception as error:
            return b''

        return data

    def close(self):
        # 写模式时，确认init写生成器创建对象，防止ftp上传空文件时没有创建对象的问题
        if 'w' in self.mode:
            self.ensure_init_write_generator()

        self._sync_cache()
        self.file.close()
        self.closed = True

    def seek(self, offset):
        self._sync_cache()      # seek前，同步可能缓存的数据
        if self.offset == offset:
            return

        self.offset = offset
        if self.offset == 0:
            self.is_breakpoint = False
        else:
            self.is_breakpoint = True

        # 当前文件指针变offset了，读、写生成器都需要根据offset从新init
        self.write_generator = None
        self.read_generator = None

    def _sync_cache(self):
        """
        缓存的文件数据同步到存储桶
        """
        lenght = self.file.tell()
        if lenght == 0:
            return

        self.ensure_init_write_generator()

        try:
            self.write_generator.send((self.offset, self.file.getvalue()))
        except HarborError as error:
            raise FilesystemError(error.msg)

        self.offset += lenght
        self.file.truncate(0)       # 清缓存
        self.file.seek(0)           # 重置缓存指针

class DownLoader(object):
    def __init__(self, bucket_name, ftp_path, client):
        self.bucket_name = bucket_name
        self.name = os.path.basename(ftp_path)
        # self.name = '/'
        self.ftp_path = ftp_path
        self.client = client
        self.closed = False
        self.id = 0
        try:
            self.obj_generator, ob = self.client.ftp_get_obj_generator(
                self.bucket_name, self.ftp_path[1:], per_size=4 * 1024 ** 2)
        except HarborError as error:
            raise FilesystemError(error.msg)

    def read(self, size=None):
        try:
            data = next(self.obj_generator)
        except Exception as error:
            return b''

        return data

    def close(self):
        self.closed = True

if __name__ == '__main__':
    print('no info')
    f = open('harbor_auto.py')
    f.read()


from .todo_utils import *

def get_pretty_time_string(delta_data):
    measure = {
        -1: "Less than minute ago",
        0: "day",
        1: "hour",
        2: "minute",
        10: "days",
        11: "hours",
        12: "minutes",
    }
    measure_value, data = get_first_not_zero_value(*delta_data)
    return (
        measure[-1] if measure_value == -1 else f"{data} {measure[measure_value]} ago"
    )

def substract_two_dates(date_given):
    diff = datetime.now() - datetime.fromisoformat(date_given)
    days = diff.days
    seconds = diff.seconds
    hours, remainder = divmod(seconds, 3600)
    minutes, seconds = divmod(remainder, 60)
    return get_pretty_time_string([days, hours, minutes])

def get_human_readable_time(date_given):
    return substract_two_dates(date_given)


import json
import ast
from keyPath import PATH as PATH

def get_Keys():
    with open(PATH) as f:
        data = json.load(f)
        parsed_data = ast.literal_eval(data)
        print("***SECRET***")
        #print(data)
        #print(parsed_data)
        NOTION_DATABASE_ID = parsed_data['NOTION_DATABASE_ID']
        NOTION_INTEGRATION_TOKEN = parsed_data["NOTION_INTEGRATION_TOKEN"]
        GOOGLE_TO_DO_DAILY_TASKLIST_ID= parsed_data["GOOGLE_TO_DO_DAILY_TASKLIST_ID"]
        return(parsed_data)
        #print(NOTION_DATABASE_ID)
        #print(NOTION_INTEGRATION_TOKEN)
        #print(GOOGLE_TO_DO_DAILY_TASKLIST_ID)

def set_Keys(dictionary_of_keys):
    keys_dict = dictionary_of_keys
    print(keys_dict)
    secret_keys_json = json.dumps(keys_dict)
    with open('secret_keys.json', 'w') as json_file:
        json.dump(secret_keys_json, json_file)

def convert_to_RFC_datetime(year=1900, month=1, day=1, hour=0, minute=0):
    dt = datetime.datetime(year, month, day, hour, minute, 0, 000).isoformat() + 'Z'
    return dt

#set_Keys({"NOTION_DATABASE_ID" : "", "NOTION_INTEGRATION_TOKEN" : "", "GOOGLE_TO_DO_DAILY_TASKLIST_ID" : ""})
#get_Keys()


from unittest.mock import sentinel

import pytest

from commando.core import commando

class TestCommando:
    def test_it(self):
        exit_status = commando(
            [sentinel.repo_1, sentinel.repo_2],
            ["test", "command"],
            sentinel.commit_message,
            sentinel.branch,
            sentinel.pr_title,
            sentinel.pr_body,
        )

        assert not exit_status

    @pytest.fixture(autouse=True)
    def subprocess(self, mocker):
        return mocker.patch("commando.core.subprocess", autospec=True)


"""strategyidx added to Position

Revision ID: 2c9f80fe65fc
Revises: 63ef97eb1cee
Create Date: 2020-04-21 10:24:50.332582

"""
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = '2c9f80fe65fc'
down_revision = '63ef97eb1cee'
branch_labels = None
depends_on = None

def upgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.add_column('position', sa.Column('strategyidx', sa.Integer(), nullable=True))
    # ### end Alembic commands ###

def downgrade():
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column('position', 'strategyidx')
    # ### end Alembic commands ###


# -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import models, migrations

class Migration(migrations.Migration):

    dependencies = [
        ('core', '0008_auto_20150321_2108'),
    ]

    operations = [
        migrations.AlterModelOptions(
            name='country',
            options={'ordering': ['iso2'], 'verbose_name': 'Country', 'verbose_name_plural': 'Countries'},
        ),
        migrations.AlterModelOptions(
            name='grouping',
            options={'ordering': ['name'], 'verbose_name': 'Channel Group', 'verbose_name_plural': 'Channel Groups'},
        ),
        migrations.AlterModelOptions(
            name='language',
            options={'ordering': ['iso2'], 'verbose_name': 'Language', 'verbose_name_plural': 'Languages'},
        ),
        migrations.RenameField(
            model_name='channel',
            old_name='primary_language',
            new_name='language',
        ),
    ]


import os
import json

from typing import List, Dict, Any

from pindex.core import Pindex
from pindex.duplicates import find_duplicates
from pindex.util import Printer, human_readable, concat

def index(config: Pindex) -> Dict[int, List[Dict[str, Any]]]:
    all_files = {}
    total_size = 0
    printer = Printer(1)
    for root, _, files in os.walk(config.root):
        relative_path = os.path.relpath(root, config.root)
        for name in files:
            printer.print(
                f"Processed {human_readable(total_size)}   \r", end="")
            entry = {
                "name": name,
                "path": relative_path
            }
            try:
                size = os.path.getsize(os.path.join(root, name))
                total_size += size
                entry["size"] = size
                if size not in all_files.keys():
                    all_files[size] = []
                all_files[size].append(entry)
            except (FileNotFoundError, PermissionError):
                pass
    printer.print_last()
    print(f"\n{len(concat(all_files.values()))} "
          "file(s) have been indexed successfully")
    return all_files

def walk(config: Pindex):
    all_files = index(config)
    if config.find_duplicates:
        find_duplicates(config, all_files.copy())

    all_files = concat(all_files.values())

    with open(config.save, "w") as f:
        f.write(json.dumps({
            "root": config.root,
            "files": all_files
        }, indent=2))


from torch.utils.data import DataLoader
from sentence_transformers import SentenceTransformer

from transformers import AutoTokenizer, AutoModel
import torch
import warnings
warnings.filterwarnings('ignore')

def expand_label(labels,sentences):
  new_labels = [0 for i in range(len(sentences))]
  for i in labels:
    new_labels[i] = 1
  return new_labels

class BERT_BASE_ENCODER:
    def __init__(self, model="bert-base-uncased", pooling="average"):
            self.bert = AutoModel.from_pretrained(model)
            self.pool=self.cls_pooling if pooling=="cls" else self.avg_pooling
            self.tokenizer = AutoTokenizer.from_pretrained(model)
            self.model = 'bert_cls_token'
            self.max_length = 512
            if torch.cuda.is_available():
                self.device = 'cuda'
                self.bert.to('cuda')
                print('Moved BERT to gpu!')
            else:
                print('No gpu is being used')
                self.device = 'cpu'

    def cls_pooling(self, model_output, attention_mask):
      return model_output[0][:,0]

    def avg_pooling(self, model_output, attention_mask):
      input_mask_expanded = attention_mask.unsqueeze(-1).expand(model_output[-1].size()).float()
      sum_embeddings = torch.sum(model_output[-1] * input_mask_expanded, 1)

      sum_mask = input_mask_expanded.sum(1)

      sum_mask = torch.clamp(sum_mask, min=1e-9)

      return sum_embeddings / sum_mask

    def encode(self, sentences, convert_to_tensor=False, batch_size=32):

        all_embeddings = []

        length_sorted_idx = np.argsort([-len(sen.split()) for sen in sentences])

        sentences_sorted = [sentences[idx] for idx in length_sorted_idx]

        for start_index in range(0, len(sentences), batch_size):

            sentences_batch = sentences_sorted[start_index:start_index + batch_size]

            encoded_input = self.tokenizer(sentences_batch, padding=True, truncation=True, return_tensors='pt',
                                           max_length=self.max_length)

            with torch.no_grad():

                if encoded_input['input_ids'].shape[0] > 100:
                    pass

                try:
                    model_output = self.bert(input_ids=encoded_input['input_ids'].to(self.device),
                                         attention_mask=encoded_input['attention_mask'].to(self.device),
                                         output_hidden_states=True).hidden_states
                except RuntimeError:
                    self.bert.to(self.device)
                    model_output = self.bert(input_ids=encoded_input['input_ids'].to(self.device),
                                             attention_mask=encoded_input['attention_mask'].to(self.device),
                                             output_hidden_states=True).hidden_states

                model_output = self.pool(model_output,
                                         encoded_input['attention_mask'].to(self.device)).detach().cpu()

                all_embeddings.extend(model_output)

        all_embeddings = [all_embeddings[idx] for idx in np.argsort(length_sorted_idx)]

        if convert_to_tensor:
            return torch.stack(all_embeddings)
        else:
            return np.asarray([emb.numpy() for emb in all_embeddings])

    def get_sentence_embedding_dimension(self):
        return 768

class ST_ENCODER:
  def __init__(self, model = "all-mpnet-base-v2"):
    self.model = SentenceTransformer(model)
  def encode(self, sentences, convert_to_tensor = False):
    if convert_to_tensor:
      return self.model.encode(sentences, pt=True)
    else:
      return self.model.encode(sentences)

  def get_sentence_embedding_dimension(self):
    return 768


import requests
import json

#url = "https://reports.sem-o.com/api/v1/documents/static-reports"
url= "https://reports.sem-o.com/api/v1/documents/static-reports?ReportName=Balancing%20and%20Imbalance%20Market%20Cost%20View&Date=>2019-11-10"
response = requests.get(url)
data = response.json()

listOfReports = []
#output to console
#print (data)
for item in data["items"]:
    #print(item["ResourceName"])
    listOfReports.append(item["ResourceName"])

for ReportName in listOfReports:
    #print(ReportName)
    url ="https://reports.sem-o.com/api/v1/documents/"+ReportName
    print(url)
    response= requests.get(url)
    aReport= response.json()

#other code
#save this to a file
filename = 'allreports.json'
# Writing JSON data
f =  open(filename, 'w')
json.dump(data, f, indent=4)


from arpeggio import Optional, ZeroOrMore, OneOrMore, EOF, StrMatch
from arpeggio import RegExMatch as _
from arpeggio import ParserPython
from arpeggio import PTNodeVisitor, visit_parse_tree
import sys
import copy

# From Arpeggio tutorial, prevents string from appearing in parse tree
class SuppressStr(StrMatch):
    suppress = True

# Parser Functions
def atomUnit():
    # Second option lists accepted atoms, can add more if necesary
    return [(SuppressStr('('), atomList, SuppressStr(')')),
            (Optional(['=','#']), _(r'Br|Cl|[BCFHINOPS]'), Optional(_(r'\d*')))]

def atomList():
    return ZeroOrMore(atomUnit)

def ring():
    return SuppressStr(':'), atomList, \
        SuppressStr(':'), Optional(['=','#'])

def fragment():
    return [ring, atomList]

def fragmentList():
    return fragment, ZeroOrMore(SuppressStr('-'),fragment), EOF

# Visitor class that walks parse tree
# Returns list of fragments,
# Each fragment is tuple of 2 lists,
# First list is starting branch,
# Second list is ring
# Any atomList in parenthesis is its own list
# Any atomUnit is tuple of Atom string and int tag
class AtomCloneVistor(PTNodeVisitor):
    def visit_fragmentList(self, node, children):
        return list(children)

    def visit_fragment(self, node, children):
        d = dict()
        for i in range(len(node)):
            if node[i].rule_name == "atomList":
                d[node[i].rule_name] = children[i]
            elif node[i].rule_name == "ring":
                d[node[i].rule_name], bond = children[i]
                d["ringBond"] = bond

        return d

    def visit_ring(self, node, children):
        if len(children) == 0 or children[0] in ['=','#']:
            return None
        bond = 1
        if len(children) == 2 and children[1] == '=':
            bond = 2
        elif len(children) == 2 and children[1] == '#':
            bond = 3

        return (children[0], bond)

    def visit_atomList(self, node, children):
        return list(children)

    def visit_atomUnit(self, node, children):
        # Parenthesized branch
        if node[0].rule_name == "atomList":
            return children[0]

        noBondChildren = children
        bond = 1
        tag = None
        if children[0] == '=':
            bond = 2
            noBondChildren = children[1:]
        elif children[0] == '#':
            bond = 3
            noBondChildren = children[1:]

        # Atom and tag
        if len(noBondChildren) == 2:
            d= dict()
            d["atom"] = noBondChildren[0]
            d["tag"] = int(noBondChildren[1])
            d["bond"] = bond
            return d
        # Atom without tag
        else:
            d= dict()
            d["atom"] = noBondChildren[0]
            d["bond"] = bond
            return d

def initParser():
    """
    Set parser global variable to arpeggio parser for our format
    """
    global parser
    parser = ParserPython(fragmentList)

def fragmentAdjacencyMatrix(fragment):
    """
    Takes tuple representing a fragment and returns adjacency matrix,
    label list, and tag dictionary.

    arguments:
    fragment - tuple representing fragment as returned from visitor class

    returns:
    adjacency matrix as list of lists, label matrix as list of strings,
    dictionary mapping tags to index
    """
    numAtoms = 0
    ringStart = None
    labels = []
    bonds = []
    tags = dict()
    fidxStack = [0]
    atomStack = [-1]
    struct = "atomList"
    if struct not in fragment:
        struct = "ring"
        ringStart = 0
    # Traverse fragment in DFS order with iterative stack
    while len(fidxStack) > 0:
        # Use struct and fidxStack to get current object
        obj = fragment[struct]
        for i in fidxStack:
            if i < len(obj):
                obj = obj[i]
            else:
                # When index at one level goes out of bounds,
                # return to previous level
                fidxStack.pop()
                # Check if finished with ring, add ring closing bond
                if len(fidxStack) == 0 and struct == "ring":
                    bonds.append((ringStart, atomStack[0], fragment["ringBond"]))
                atomStack.pop()
                if len(fidxStack) > 0:
                    fidxStack[-1] += 1
                obj = None
                break
        # Handle individual atom by adding tag, label, and implicit bond
        if isinstance(obj, dict):
            labels.append(obj["atom"])
            if "tag" in obj:
                # assert (obj["tag"] not in tags), "Duplicate tag in same fragment"
                tags[obj["tag"]] = numAtoms

            bonds.append((atomStack[-1],numAtoms,obj["bond"]))
            fidxStack[-1] += 1
            atomStack[-1] = numAtoms
            numAtoms += 1
        # Handle branch (atomList) by setting up parent atom for next bond
        elif isinstance(obj, list):
            fidxStack.append(0)
            # atomStack.append(numAtoms-1)
            atomStack.append(atomStack[-1])

    # Remove dummy starter bond
    bonds = bonds[1:]

    # Make adjacency matrix with all zeros
    mat = []
    for i in range(numAtoms):
        mat.append([0]*numAtoms)

    for (i,j, b) in bonds:
        mat[i][j] = b
        mat[j][i] = b

    return mat, labels, tags

def mergeFragments(a, b):
    """
    Merges two fragment adjacency matrices and labels using tags.

    arguments:
    a - Fragment tuple (matrix, label, tag)
    b - Second fragment tuple (matrix, label, tag)

    returns:
    Merged fragment tuple (matrix, label, tag)
    """
    matA = a[0]
    labA = a[1]
    tagA = a[2]

    matB = b[0]
    labB = b[1]
    tagB = b[2]

    # Compute final amount of atoms in mergeed fragment
    newNumAtom = len(labA) + len(labB)
    intersect = [k for k in tagB if k in tagA]
    newNumAtom -= len(intersect)

    newLab = copy.copy(labA)
    # Construct list mapping index in B to combined index (accounting for tags)
    idxB = [None] * len(labB)
    for k in intersect:
        assert labA[tagA[k]] == labB[tagB[k]], "Same tag on different atom types"
        idxB[tagB[k]] = tagA[k]

    currIdx = len(labA)
    for i in range(len(idxB)):
        if idxB[i] is None:
            idxB[i] = currIdx
            newLab.append(labB[i])
            currIdx += 1
    for k in tagB.keys():
        tagB[k] = idxB[tagB[k]]
    newTag = {**tagB, **tagA}

    newMat = []
    for i in range(newNumAtom):
        newMat.append([0]*newNumAtom)

    # Copy A matrix
    for i in range(len(matA)):
        for j in range(len(matA)):
            newMat[i][j] = matA[i][j]

    # Add bonds from B matrix
    for i in range(len(matB)):
        for j in range(len(matB)):
            newMat[idxB[i]][idxB[j]] = max(matB[i][j],
                    newMat[idxB[i]][idxB[j]])

    return (newMat, newLab, newTag)

def strToAdjacencyMatrix(s):
    """
    Parses string and returns adjacency matrix and label array

    arguments:
    s - string to be parsed

    returns:
    adjacency matrix as list of lists, label matrix as list of strings
    """
    initParser()
    tree = parser.parse(s)
    listForm = visit_parse_tree(tree, AtomCloneVistor())
    currFragment = fragmentAdjacencyMatrix(listForm[0])
    for nextList in listForm[1:]:
        nextFragment = fragmentAdjacencyMatrix(nextList)
        currFragment = mergeFragments(currFragment, nextFragment)

    return currFragment[0:2]

def printMatrix(mat):
    """
    Prints matrix (list of lists) in nice format
    (assumes all entries take same number of characters)

    arguments:
    mat - list of lists representing matrix
    """
    for row in mat:
        print(' '.join([str(x) for x in row]))

def printLabels(lab):
    """
    Prints label list in nice format

    arguments:
    mat - list of lists representing matrix
    """
    for i,c in enumerate(lab):
        print(str(i)+" "+c)

# initParser()
# test = ":C(C)C(C)C0C1S:-:C(CNC2)CCCO:-:C2C0C1NCN:"
# mat,lab = strToAdjacencyMatrix(test)
# print(test)
# printMatrix(mat)
# printLabels(lab)
#assert(len(sys.argv) == 2)

# initParser()

# with open(sys.argv[1]) as f:
#     for line in f:
#         cleanLine = line.strip()
#         print("String: ", cleanLine)
#         matrix, lab = strToAdjacencyMatrix(cleanLine)
#         printMatrix(matrix)
#         printLabels(lab)