from ast import parse, walk, FunctionDef, ClassDef, Name
from re import findall

from math import log, sqrt

from pandas import DataFrame
from nltk.stem import PorterStemmer
from sklearn.metrics.pairwise import cosine_similarity

from .Analyzer import Analyzer

class SemanticAnalyzer(Analyzer):

    def __init__(self, url, remove_bulk=-1) -> None:
        super().__init__(url, remove_bulk)

        self.couplings_type = 'semantic'

        self.tf_idf_df = None
        self.file_to_identifiers = None

    def compute_couplings(self):

        self.file_to_identifiers = self.get_corpus()

        self.preprocess_words(self.file_to_identifiers)

        voc_size, voc_to_index = self.compute_voc(self.file_to_identifiers)

        tf = self.compute_tf(voc_to_index, self.file_to_identifiers)
        idf = self.compute_idf(voc_to_index, self.file_to_identifiers)
        tf_idf = self.compute_tf_idf(voc_to_index, tf, idf)

        self.tf_idf_df = DataFrame.from_dict(tf_idf, orient='index')

        self.run_general_analysis(
            get_logical_couplings_df=True,
            get_commit_to_files=True,
            get_dates=True
            )

    def get_distance_matrix(self):

        distance_matrix = cosine_similarity(self.tf_idf_df)
        for i in range(len(distance_matrix)):
            distance_matrix[i][i] = 1
        distance_df = DataFrame(distance_matrix, index=self.tf_idf_df.index, columns=self.tf_idf_df.index)

        self.distance_matrix = distance_df

        return distance_df

    def get_corpus(self):
        """ Get a list of identifiers of each file in a repo.
        """

        file_to_identifiers = {}
        for file_path in self.repo_files_path:

            try :
                with open(self.repo_folder + '\\' + file_path) as data_source:
                    ast_root = parse(data_source.read())

                identifiers = []

                for node in walk(ast_root):
                    if isinstance(node, Name):
                        identifiers.append(node.id)
                    elif isinstance(node, FunctionDef) or isinstance(node, ClassDef):
                        identifiers.append(node.name)

                file_to_identifiers[file_path] = identifiers

            except:
                path = self.repo_folder + "\\" + file_path
                print(f'Could not read {path}')
                pass

        return file_to_identifiers

    def preprocess_words(self, file_to_identifiers):
        """ Preprocess words for further analysis : stems, split and lower words.
        """

        stemmer = PorterStemmer()

        for key in file_to_identifiers.keys():

            old_words = file_to_identifiers[key]
            new_words = []

            for i in range(len(old_words)):

                splitted_sentence = self.split_sentence(old_words[i])
                for word in splitted_sentence:
                    new_words.append(word)

            new_words = [str.lower(word) for word in new_words]
            new_words = [stemmer.stem(word) for word in new_words]

            file_to_identifiers[key] = new_words

    @staticmethod
    def split_sentence(word):
        """ Split a snake or camel case string into its composing words.
        """

        # Snake split
        splitted_snake_sentence = word.split('_')

        # camel_word = re.sub(r'_(.)', lambda m: m.group(1).upper(), word)

        splitted_sentence = []
        for snake_word in splitted_snake_sentence:
            camel_words = findall(r'.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)', snake_word)
            for camel_word in camel_words:
                splitted_sentence.append(camel_word)

        return splitted_sentence

    @staticmethod
    def compute_voc(file_to_identifiers):
        """ Compute the vocabulary of a repo (list of all the words).
        """

        voc = set()

        for _, value in file_to_identifiers.items():

            for word in value:

                if word not in voc:
                    voc.add(word)

        voc_size = len(voc)

        voc_to_index = {}

        i = 0
        for word in voc:
            voc_to_index[word] = i
            i += 1

        return voc_size, voc_to_index

    @staticmethod
    def compute_tf(voc_to_index, file_to_identifiers):
        """ Compute the tf values for each file in repo.
        """

        tf = {}

        for file_path in file_to_identifiers.keys():

            tf[file_path] = [0 for _ in range(len(voc_to_index))]

            for word in file_to_identifiers[file_path]:

                tf[file_path][voc_to_index[word]] += 1

            num_identifiers = len(file_to_identifiers[file_path])

            if num_identifiers > 0:
                tf[file_path] = [value / num_identifiers for value in tf[file_path]]

        return tf

    @staticmethod
    def compute_idf(voc_to_index, file_identifiers):
        """ Compute the idf values for each file in repo.
        """

        idf = {}

        for word in voc_to_index.keys():

            num_doc = 0

            for identifiers in file_identifiers.values():

                if word in identifiers:
                    num_doc += 1

            idf[word] = log(len(file_identifiers) / num_doc)

        return idf

    @staticmethod
    def compute_tf_idf(voc_to_index, tf, idf):
        """ Compute the tf_idf values for each file in repo.
        """

        tf_idf = {}

        for file_path in tf.keys():

            tf_idf[file_path] = [0 for _ in range(len(voc_to_index))]

            for word, index in voc_to_index.items():

                tf_idf[file_path][index] = tf[file_path][index] * idf[word]

        return tf_idf

    @staticmethod
    def compute_cosine_distance(a, b):
        """ Compute the cosine distance between two vectors a and b.
        """

        norm_a = 0
        norm_b = 0

        dot = 0

        for i in range(len(a)):

            dot += a[i] * b[i]

            norm_a += a[i] ** 2
            norm_b += b[i] ** 2

        norm_a = sqrt(norm_a)
        norm_b = sqrt(norm_b)

        return dot / (norm_a * norm_b)


# pylint: disable=invalid-name
from typing import Callable, TypeVar, Any, Optional

from .AbstractStateMachineFieldMapper import AbstractStateMachineFieldMapper
from .MessageFormatter import MessageFormatter
from .SingleLineMesasgeFormatter import SingleLineMesasgeFormatter
from .statemachine.CodeError import CodeError
from .statemachine.Defaulter import Defaulter
from .statemachine.End import End
from .statemachine.Getter import Getter
from .statemachine.MachineContext import MachineContext
from .statemachine.Setter import Setter
from .statemachine.State import State
from .statemachine.Translator import Translator
from .statemachine.Validator import Validator
from .statemachine.WarnIfDefinedOrDataError import WarnIfDefinedOrDataError
from ..Field import Field
from ..MappingContext import MappingContext
from ...IllegalStateException import IllegalStateException

OO = TypeVar("OO")
RO = TypeVar("RO")
P = TypeVar("P")

class MandatoryFieldMapper(AbstractStateMachineFieldMapper[OO, RO, P]):
    """State machine based mapper for mandatory fields.

    Mandatory means that corresponding field in the resulting object must be set.
    However, in pure validation mode when there is no field to set in
    resulting object, mandatory means that input value must be present and must be valid.

    It is considered to be data error if mandatory field is absent in input
    or has an invalid value (as per validator) or can't be translated.
    Data error means that an exception will be raised.

    If there is default value then it will be used if input field is missing,
    invalid or causes translation error. However, warning message will still
    be generated.

    If either getter not defaulter are set - no way to obtain value to validate/translate/set.
    This is treated as code error.

    If we have a getter then there supposed to be either validator to check input value or
    setter to propagate it to output. If there are none then the field should not be mandatory.

    Absence means null.

    OO - Original object type.

    RO - Resulting object type.

    P - Type of parameters object.
    """

    class WarnIfDefinedOrDataError_InnerClass(WarnIfDefinedOrDataError):
        def __init__(self, defaulter, messageFormatter, text):
            super().__init__(defaulter, messageFormatter)
            self.messageFormatter = messageFormatter
            self.text = text

        def createMessage(
                self, field: Field, ctx: MappingContext, machine: MachineContext
        ):
            return self.messageFormatter.formatMessage(
                ctx, field, self.text.format(**locals())
            )

    def __init__(
        self, messageFormatterFactory: Optional[Callable[..., MessageFormatter]] = None
    ):
        messageFormatter: MessageFormatter = (
            messageFormatterFactory()
            if messageFormatterFactory
            else SingleLineMesasgeFormatter(mandatory=True)
        )
        """
        Creates new instance.

        :param messageFormatterFactory: Factory for formatter of error messages or null
                                         if {@link SingleLineMesasgeFormatter} should be used.
        """
        getter = Getter()
        validator = Validator()
        translator = Translator()
        defaulter = Defaulter()
        setter = Setter()
        end = End()

        getter.configure(
            validator,
            MandatoryFieldMapper.WarnIfDefinedOrDataError_InnerClass(
                defaulter, messageFormatter, "Input value absent"
            ),
            MandatoryFieldMapper.WarnIfDefinedOrDataError_InnerClass(
                defaulter, messageFormatter, "Cannot obtain input value"
            ),
            defaulter,
        )
        validator.configure(
            translator,
            MandatoryFieldMapper.WarnIfDefinedOrDataError_InnerClass(
                defaulter,
                messageFormatter,
                "Input value is invalid: '{machine.resultValue}'",
            ),
            MandatoryFieldMapper.WarnIfDefinedOrDataError_InnerClass(
                defaulter,
                messageFormatter,
                "Input value cannot be validated: '{machine.resultValue}'",
            ),
            translator,
        )
        translator.configure(
            setter,
            CodeError("Translator returns null", messageFormatter),
            MandatoryFieldMapper.WarnIfDefinedOrDataError_InnerClass(
                defaulter,
                messageFormatter,
                "Input value cannot be mapped: '{machine.resultValue}'",
            ),
            setter,
        )

        defaulter.configure(
            setter,
            CodeError("Defaulter returns null, this makes no sense", messageFormatter),
            CodeError("Defaulter throws exception", messageFormatter),
            CodeError(
                "Should never see this as should get data error instead"
                + " by the means of WarnIfDefinedOrDataError.simpleName",
                messageFormatter,
            ),
        )
        setter.configure(
            end, end, CodeError("Setter throws exception", messageFormatter), end
        )
        self.stateMachine = getter

    def getStateMachine(self, field: Field[OO, RO, Any, Any, P]) -> State:
        assert field is not None
        if not field.getter and not field.defaulter:
            raise IllegalStateException(
                "Neither getter not defaulter are set"
                + " - no way to obtain value to validate/translate/set."
            )

        if field.getter and not field.validator and not field.setter:
            raise IllegalStateException(
                "There is a getter but neither validator, nor setter"
                + " so we neither validate, nor propagate it - this makes no sense for mandatory field."
            )

        return self.stateMachine


import os
import tempfile
import unittest
import torch
from torch import nn

from towhee.trainer.optimization.adafactor import Adafactor
from towhee.trainer.optimization.adamw import AdamW
from towhee.trainer.scheduler import (
    configure_constant_scheduler,
    configure_constant_scheduler_with_warmup,
    configure_linear_scheduler_with_warmup,
    configure_cosine_scheduler_with_warmup,
    configure_cosine_with_hard_restarts_scheduler_with_warmup,
    configure_polynomial_decay_scheduler_with_warmup
)

def unwrap_scheduler(scheduler, num_steps=10):
    lr_sch = []
    for _ in range(num_steps):
        lr_sch.append(scheduler.get_lr()[0])
        scheduler.step()
    return lr_sch

def unwrap_and_save_reload_scheduler(scheduler, num_steps=10):
    lr_sch = []
    for step in range(num_steps):
        lr_sch.append(scheduler.get_lr()[0])
        scheduler.step()
        if step == num_steps // 2:
            with tempfile.TemporaryDirectory() as tmpdirname:
                file_name = os.path.join(tmpdirname, 'schedule.bin')
                torch.save(scheduler.state_dict(), file_name)

                state_dict = torch.load(file_name)
                scheduler.load_state_dict(state_dict)
    return lr_sch

class ScheduleInitTest(unittest.TestCase):
    mdl = nn.Linear(50, 50)
    optimizer = AdamW(mdl.parameters(), lr=10.0)
    num_steps = 10

    def assertListAlmostEqual(self, list1, list2, tol, msg=None):
        self.assertEqual(len(list1), len(list2))
        for a, b in zip(list1, list2):
            self.assertAlmostEqual(a, b, delta=tol, msg=msg)

    def test_schedulers(self):
        common_kwargs = {'num_warmup_steps': 2, 'num_training_steps': 10}
        # schedulers doct format
        # function: (sched_args_dict, expected_learning_rates)
        scheds = {
            configure_constant_scheduler: ({}, [10.0] * self.num_steps),
            configure_constant_scheduler_with_warmup: (
                {'num_warmup_steps': 4},
                [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
            ),
            configure_linear_scheduler_with_warmup: (
                {**common_kwargs},
                [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
            ),
            configure_cosine_scheduler_with_warmup: (
                {**common_kwargs},
                [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
            ),
            configure_cosine_with_hard_restarts_scheduler_with_warmup: (
                {**common_kwargs, 'num_cycles': 2},
                [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
            ),
            configure_polynomial_decay_scheduler_with_warmup: (
                {**common_kwargs, 'power': 2.0, 'lr_end': 1e-7},
                [0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
            ),
        }

        for scheduler_func, data in scheds.items():
            kwargs, expected_learning_rates = data

            scheduler = scheduler_func(self.optimizer, **kwargs)
            self.assertEqual(len([scheduler.get_lr()[0]]), 1)
            lr_sch_1 = unwrap_scheduler(scheduler, self.num_steps)
            self.assertListAlmostEqual(
                lr_sch_1,
                expected_learning_rates,
                tol=1e-2,
                msg=f'failed for {scheduler_func} in normal scheduler',
            )

            scheduler = scheduler_func(self.optimizer, **kwargs)
            lr_sch_2 = unwrap_and_save_reload_scheduler(scheduler, self.num_steps)
            self.assertListEqual(lr_sch_1, lr_sch_2, msg=f'failed for {scheduler_func} in save and reload')

class OptimizationTest(unittest.TestCase):
    def assertListAlmostEqual(self, list1, list2, tol):
        self.assertEqual(len(list1), len(list2))
        for a, b in zip(list1, list2):
            self.assertAlmostEqual(a, b, delta=tol)

    def test_adam_w(self):
        w = torch.tensor([0.1, -0.2, -0.1], requires_grad=True)
        target = torch.tensor([0.4, 0.2, -0.5])
        criterion = nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        optimizer = AdamW(params=[w], lr=2e-1, weight_decay=0.0)
        for _ in range(100):
            loss = criterion(w, target)
            loss.backward()
            optimizer.step()
            w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
            w.grad.zero_()
        self.assertListAlmostEqual(w.tolist(), [0.4, 0.2, -0.5], tol=1e-2)

    def test_adafactor(self):
        w = torch.tensor([0.1, -0.2, -0.1], requires_grad=True)
        target = torch.tensor([0.4, 0.2, -0.5])
        criterion = nn.MSELoss()
        # No warmup, constant schedule, no gradient clipping
        optimizer = Adafactor(
            params=[w],
            lr=1e-2,
            eps=(1e-30, 1e-3),
            clip_threshold=1.0,
            decay_rate=-0.8,
            beta1=None,
            weight_decay=0.0,
            relative_step=False,
            scale_parameter=False,
            warmup_init=False,
        )
        for _ in range(1000):
            loss = criterion(w, target)
            loss.backward()
            optimizer.step()
            w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
            w.grad.zero_()
        self.assertListAlmostEqual(w.tolist(), [0.4, 0.2, -0.5], tol=1e-2)

if __name__ == '__main__':
    unittest.main()


from netmiko import ConnectHandler
from concurrent.futures import ThreadPoolExecutor
from my_devices import all_devices
from datetime import datetime

def connector(device, command):
    with ConnectHandler(**device) as connect:
        output = connect.send_command(command)
        return output

if __name__ == "__main__":

    start_time = datetime.now()

    for device in all_devices:
        optional_args = device.pop("optional_args")
        print(f"#### Command started for {device['host']} at {datetime.now()}")
        output = connector(device, "show version")
        print(output)
        print(f"#### Command finished for {device['host']} at {datetime.now()}\n")

    print(f"#### Job was completed in {str(datetime.now() - start_time)}")