"""This scripts includes DTI registration pipeline using DTI-TK tool.
I followed the official DTI-TK tutorials.
The steps listed below are taken from here:
http://dti-tk.sourceforge.net/pmwiki/pmwiki.php?n=Documentation.Registration
"""

import os
import shutil
import logging
import subprocess
from typing import Union
from pathlib import Path
from dtip.convert import fsl_to_dtitk_multi
from dtip.utils import show_exec_time

# * Add dtitk tool to PATH
from dtip.utils import ROOT_DIR
if os.getenv('DTITK_ROOT') is None:
    # Add DTI-TK PATH as environment variable
    dtitk_maindir = f"{ROOT_DIR}/dtitk"
    os.environ["DTITK_ROOT"] = dtitk_maindir
    os.environ["PATH"] += f":{dtitk_maindir}/bin:{dtitk_maindir}/utilities:{dtitk_maindir}/scripts"

__all__ = [
    'dtitk_register_multi',
    'template_to_subject_space',
    'template_to_subject_space_multi'
]

@show_exec_time
def dtitk_register_multi(
    input_path: Union[str, Path],
    template_path: Union[None, str, Path],
    bootstrapped_template_path: Union[None, str, Path],
    output_path: Union[str, Path],
    n_diffeo_iters: int = 8,
) -> int:
    """DTI existing Template-based Image Registration using Diffusion
        Tensor Imaging ToolKit (DTI-TK)

    Args:
        input_path: folder path containing a subject's data.
        template_path: Path of the template to use for registration.
        output_path: location to save the output files.

    Returns:
        exit code 0 on successful execution.
    """

    input_path, output_path = Path(input_path), Path(output_path)
    if bootstrapped_template_path:
        bootstrapped_template_path = Path(bootstrapped_template_path)

    if bootstrapped_template_path is None:
        # * Step 1: Convert FSL format to DTI-TK format and move files
        # * to `output_path`.
        exit_code = fsl_to_dtitk_multi(input_path, output_path)
        if exit_code != 0:  # Stop here if any error
            _msg = "Error in `dtitk_register_multi` execution :(. Stopped."
            logging.error(_msg)
            raise RuntimeError(_msg)
        logging.info(
            f"Converted `{input_path}` to DTI-TK format and saved at `{output_path}`."
        )

        # * SPATIAL NORMALIZATION AND ATLAS CONSTRUCTION
        # * Step 2. Bootstrapping the initial DTI template from the input DTI volumes
        # Get subjects' DTI file paths
        subs_filepaths = [
            subject_path / "dti_dtitk.nii.gz"
            for subject_path in output_path.glob("*")
            if subject_path.is_dir()
        ]
        # Create a file with subset names
        subs_filepath = output_path / "subs.txt"
        with open(subs_filepath, "w") as subf:
            for filepath in subs_filepaths:
                subf.write(f"{filepath}\n")

        print(subs_filepath)
        # Run the `dti_template_bootstrap` command
        logging.info("Starting template bootstrapping...")
        subprocess.run(
            [
                "dti_template_bootstrap",
                template_path,
                subs_filepath,
                "EDS",
                "4",
                "4",
                "4",
                "0.0001",
            ]
        )
        logging.info("Bootstrapping complete.")
    else:

        subs_filepaths = [
            subject_path / "dti_dtitk.nii.gz"
            for subject_path in output_path.glob("*")
            if subject_path.is_dir()
        ]
        # Create a file with subset names
        subs_filepath = output_path / "subs.txt"
        with open(subs_filepath, "w") as subf:
            for filepath in subs_filepaths:
                subf.write(f"{filepath}\n")

        template_path = bootstrapped_template_path

    # * Step 3: Rigid Alignment of DTI Volumes
    # As mentioned in tutorial, rigid alignment is not required when using
    # an existing template.

    # * Step 4: Affine Alignment of DTI volumes
    logging.info("Affine aligment with template refinement")
    for i, subject_path in enumerate(subs_filepaths, start=1):
        logging.info(
            f"[{i}/{len(subs_filepaths)}] Affine alignment of `{subject_path}`..."
        )
        subprocess.run(
            [
                "dti_affine_reg",
                template_path,
                subject_path,
                "EDS",
                "4",
                "4",
                "4",
                "0.001",
            ]
        )
    logging.info("Affine alignment completed!")

    # * Step 5: Deformable alignment with template refinement
    # generate the mask image for mask argument of diffeomorphic command
    logging.debug("Extract TR image from template to create a mask...")
    ret_code = subprocess.run(
        ["TVtool", "-in", template_path, "-tr", "-out", "template_tr.nii.gz"]
    ).returncode
    if ret_code != 0:
        raise RuntimeError(
            """[Error @ `dtitk_register_multi`]
        Problem with mask generation for diffeomorphic alignment."""
        )
    logging.debug("Generating mask for diffeomorphic alignment...")
    ret_code = subprocess.run(
        [
            "BinaryThresholdImageFilter",
            "template_tr.nii.gz",
            "mask.nii.gz",
            "0.01",
            "100",
            "1",
            "0",
        ]
    ).returncode
    if ret_code != 0:
        raise RuntimeError(
            """[Error @ `dtitk_register_multi`]
        Problem with BinaryThresholdImageFilter command in subprocess."""
        )
    logging.info("Done!")
    # Run alignment subject wise
    # Get subjects' DTI file paths
    for subject_path in output_path.glob("*"):
        if subject_path.is_dir():
            logging.info(
                f"Diffeomorphic alignment of subject `{subject_path}`...")
            filepath = f"{subject_path}/dti_dtitk_aff.nii.gz"
            # Run alignment
            subprocess.run(
                [
                    "dti_diffeomorphic_reg",
                    template_path,
                    filepath,
                    "mask.nii.gz",
                    "1",
                    str(n_diffeo_iters),
                    "0.002",
                ]
            )
            logging.info("Subject alignment complete!")

    # Move extra generated files to output_path folder
    logging.info(f"Moving generated files to `{output_path}`...")
    if Path('mean_initial.nii.gz').exists():
        shutil.move('mean_initial.nii.gz', output_path/'mean_initial.nii.gz')
    if Path('mask.nii.gz').exists():
        shutil.move('mask.nii.gz', output_path/'mask.nii.gz')
    if Path('template_tr.nii.gz').exists():
        shutil.move('template_tr.nii.gz', output_path/'template_tr.nii.gz')
    if Path('bootstrapped_template.nii.gz').exists():
        shutil.move('bootstrapped_template.nii.gz',
                    output_path/'bootstrapped_template.nii.gz')
    logging.info("Registration complete!")
    return 0

def template_to_subject_space(subject_dir_path: Union[str, Path],
                              template_path: Union[str, Path],
                              transform_type: str) -> int:
    """Transform the template to the subject space.
        This function uses DTI-TK tool to perform transformations.

        This function is based on the following tutorial:
            http://dti-tk.sourceforge.net/pmwiki/pmwiki.php?n=Documentation.OptionspostReg

    Args:
        subject_dir_path: Subject folder path containing `subj.nii.gz` and
            `subj.aff` files to perform affine or diffeomorphic transformation.
            If the `transform_type` argument is `diffeo` then `subj.df.nii.gz`
            file should be present as well.
        template_path: Path to the template or atlas file to transform to
            the subject space.
        transform_type: Either `affine` for affine transformation or `diffeo`
            for diffeomorphic transformation.

    Returns:
        exit code 0 on successful execution
    """

    subject_dir_path = Path(subject_dir_path)
    if not subject_dir_path.is_dir():
        raise ValueError("`subject_dir_path` must be a folder.")
    subject_path = subject_dir_path/'dti_dtitk.nii.gz'
    aff_path = subject_dir_path/'dti_dtitk.aff'
    inv_aff_path = subject_dir_path/'dti_dtitk_inv.aff'
    if transform_type == 'affine':  # perform affine transformation
        savepath = subject_dir_path / \
            Path(template_path).stem.replace('.nii', '')
        savepath = f"{savepath}_affine_dti_space.nii.gz"
        # Compute the inverse of affine matrix
        ret_code = subprocess.run([
            'affine3Dtool', '-in', aff_path, '-invert', '-out', inv_aff_path
        ]).returncode
        if ret_code != 0:
            raise RuntimeError("Something wrong with `affine3Dtool` subprocess.")

        ret_code = subprocess.run([
            'affineScalarVolume', '-in', template_path, '-trans', inv_aff_path,
            '-target', subject_path, '-interp', '1', '-out', savepath
        ]).returncode

        if ret_code != 0:
            raise RuntimeError(
                "Something wrong with `affineScalarVolume` subprocess.")
    elif transform_type == 'diffeo':  # perform diffeomorphic transformation
        savepath = subject_dir_path / \
            Path(template_path).stem.replace('.nii', '')
        savepath = f"{savepath}_diffeo_dti_space.nii.gz"
        #* Combined displacement field from the native space to the template space
        diffeo_df_path = subject_dir_path/'dti_dtitk_aff_diffeo.df.nii.gz'
        combined_df_path = subject_dir_path/'dti_dtitk_combined.df.nii.gz'
        ret_code = subprocess.run([
            'dfRightComposeAffine',
            '-aff', aff_path,
            '-df', diffeo_df_path,
            '-out', combined_df_path
        ]).returncode
        if ret_code != 0:
            raise RuntimeError("Something wrong with `dfRightComposeAffine` subprocess.")
        #* Mapping the subject data to the template space
        subject_combined_path = subject_dir_path/'dti_dtitk_combined.nii.gz'
        ret_code = subprocess.run([
            'deformationSymTensor3DVolume',
            '-in', subject_path,
            '-trans', combined_df_path,
            '-target', template_path,
            '-out', subject_combined_path
        ]).returncode
        if ret_code != 0:
            raise RuntimeError("Something wrong with `dfRightComposeAffine` subprocess.")
        #* Combined displacement field from the template space to the native space
        ## 1. compute the inverse of the affine transformation with the
        ## command affine3Dtool
        ret_code = subprocess.run([
            'affine3Dtool', '-in', aff_path, '-invert', '-out', inv_aff_path
        ]).returncode
        if ret_code != 0:
            raise RuntimeError("Something wrong with `affine3Dtool` subprocess.")
        ## 2. compute the inverse of the deformable transformation with the
        ## command dfToInverse
        diffeo_inv_df_path = subject_dir_path/'dti_dtitk_aff_diffeo.df_inv.nii.gz'
        ret_code = subprocess.run([
            'dfToInverse', '-in', diffeo_df_path, '-out', diffeo_inv_df_path
        ]).returncode
        if ret_code != 0:
            raise RuntimeError("Something wrong with `dfToInverse` subprocess.")
        ## 3. combine the two inverted transformations with the command
        ## dfLeftComposeAffine (not dfRightComposeAffine)
        combined_inv_df_path = subject_dir_path/'dti_dtitk_combined.df_inv.nii.gz'
        ret_code = subprocess.run([
            'dfLeftComposeAffine',
            '-df', diffeo_inv_df_path,
            '-aff', inv_aff_path,
            '-out', combined_inv_df_path
        ]).returncode
        if ret_code != 0:
            raise RuntimeError("Something wrong with `dfLeftComposeAffine` subprocess.")
        #* Mapping the atlas space data to the subject space
        ret_code = subprocess.run([
            'deformationScalarVolume',
            '-in', template_path,
            '-trans', combined_inv_df_path,
            '-target', subject_path,
            '-interp', '1',
            '-out', savepath
        ]).returncode
        if ret_code != 0:
            raise RuntimeError("Something wrong with `deformationScalarVolume` subprocess.")
    else:
        raise NotImplementedError("Unknown `transform_type`. Only `affine` and `diffeo` is supported.")
    return 0

@show_exec_time
def template_to_subject_space_multi(subjects_dir_path: Union[str, Path],
                                    template_path: Union[str, Path],
                                    transform_type: str) -> int:
    """Transform the template to the subject space.
        This function uses DTI-TK tool to perform transformations.

        This function is based on the following tutorial:
            http://dti-tk.sourceforge.net/pmwiki/pmwiki.php?n=Documentation.OptionspostReg

    Args:
        subjects_dir_path: Parent folder path containing subject folder
            containing `subj.nii.gz` and `subj.aff` files to perform
            affine or diffeomorphic transformation. If the `transform_type`
            argument is `diffeo` then `subj.df.nii.gz` file should be present
            as well.
        template_path: Path to the template or atlas file to transform to
            the subject space.
        transform_type: Either `affine` for affine transformation or `diffeo`
            for diffeomorphic transformation.

    Returns:
        exit code 0 on successful execution
    """
    error_list = []
    template_path = Path(template_path)
    subjects_paths = [
        p for p in Path(subjects_dir_path).glob('*')
        if p.is_dir()
    ]
    total_subjects = len(subjects_paths)
    logging.debug("Starting template to subject space transform...")
    for i, subject_path in enumerate(subjects_paths, start=1):
        ret_code = template_to_subject_space(subject_dir_path=subject_path,
                                             template_path=template_path,
                                             transform_type=transform_type)
        if ret_code == 0:
            logging.info(
                f"[{i}/{total_subjects}] Transformed `{template_path.name}` to `{subject_path.name}`")
        else:
            logging.warning(
                f"Error transforming {subject_path.stem}. Skipped.")
            error_list.append(subject_path.stem)

    if error_list:
        print("=" * 10, f"{len(error_list)} Subjects with Errors", "=" * 10)
        for es in error_list:
            print(es)
        print("=" * 44)
        return len(error_list)
    else:
        print("=" * 30)
        print("All subjects completed without errors.")
        return 0


# -*- coding: utf-8 -*-
""" Tools used by Shiba data retrieving classes"""
import re
from collections import OrderedDict

import requests
import xmltodict

from furl import furl
from lxml import etree
from lxml import objectify
from lxml.etree import XMLSyntaxError

from ftfy import fix_text

from .shibaconnection import ShibaConnection
from .shibaresponseobject import ShibaResponseObject
from .shibaexceptions import (
    ShibaParameterError,
    ShibaLoginError,
    ShibaQuotaExceededError,
    ShibaRightsError,
    ShibaConnectionError,
    ShibaUnknownServiceError,
    ShibaServiceError,
    ShibaCallingError,
)

from http.client import HTTPException

def post_request(url, data):
    """Method creating and submitting the multipart request of the wished to be imported XML file.

    :param data: raw XML as text

    :rtype: plain text from servers response
    """
    header = {
        "User-Agent": "Mozilla/5.0 (Windows; U; Windows NT 6.1; de-DE; rv:1.9.0.10) "
        "Gecko/2009042316 Firefox/3.0.10 (.NET CLR 4.0.20506)"
    }
    d = {"file": data}
    r = requests.post(url, files=d, headers=header)
    return r.text

def retrieve_obj_from_url(url, data=None):
    """Give it an URL, will send you back an object based on received XML from WebService, it removes and store
    namespace as its not really useful for common usages of this API.
    You can also give it a dict or an objectified tree to POST a file (used for import webservice)
    :param url: fully formatted URL from url_constructor
    :param data: if not None, will send a POST request to URL instead of a GET one, with multipart sending of
    this parameter content

    :rtype : ShibaResponseObject object, containing lxml.objectify class, raw XML as Unicode string and the
    namespace from the XML.
    """
    try:
        if data is not None:
            xml = post_request(url, data)
        else:
            xml = requests.get(url).text
    except requests.ConnectionError:
        raise ShibaConnectionError("HTTP error = Connection error - On URL: " + url)
    except requests.HTTPError as e:
        raise ShibaConnectionError("URL error = " + e.reason + " - On URL: " + url)
    except HTTPException:
        raise ShibaConnectionError("HTTP unknown error =" + " - On URL: " + url)

    namespace = re.search(pattern='xmlns="[^"]', string=xml)
    if namespace is not None:
        namespace = namespace.group()
    else:
        namespace = ""

    xmlepured = re.sub(pattern=' xmlns="[^"]+"', repl="", string=xml, flags=0)
    xmlepured = fix_text(xmlepured).encode("utf-8")

    try:
        obj = objectify.fromstring(xmlepured)
    except XMLSyntaxError:
        print(xmlepured)
        raise

    if _check_errors(obj) is not False:
        try:
            if "Unknown error" == obj.error.code:
                raise ShibaServiceError(
                    "Unknown error from WebService (maybe the sale isn't confirmed yet?)"
                    " : "
                    + obj.error.message
                    + " - Reason : "
                    + obj.error.details.detail
                )
        except ShibaServiceError:
            raise ShibaServiceError(
                "Unknown error from WebService (maybe the sale isn't confirmed yet?)"
                " : " + obj.error.message + " - Reason : " + obj.error.details.detail
            )
        except Exception:
            raise ShibaUnknownServiceError(
                "An unknown error from the WebService has occurred - XML dump : "
                + etree.tostring(obj)
            )
    return ShibaResponseObject(namespace, obj, xml.encode("utf-8"))

def create_xml_from_item_obj(inv):
    """Generate XML from the "inv" parameter.

    :param inv: is an object or a dict hierarchized as the XML structure described \
    in the WebServices documentation. Take a look at the xmltodict or the lxml.objectify documentation too.

    :rtype: string output from objectified element or ElementTree element tostring() methods
    """
    if type(inv) is etree._ElementTree:
        return etree.tostring(inv)
    elif type(inv) is dict or type(inv) is OrderedDict:
        try:
            return xmltodict.unparse(inv)
        except Exception:
            raise ShibaCallingError(
                "error from dictionary to xml : can't create XML from dictionary,"
                " refers to xmltodict documentation"
                "(you don't need to add a root element 'items' to your items)"
            )
    else:
        raise ShibaCallingError(
            "error : bad input parameter given, "
            "expecting dict, objectify object or ElementTree element"
        )

def inf_constructor(shibaconnection, action, **kwargs):
    """Just a selective concatenation of args given and related static data from the given action.

    :param shibaconnection: ShibaConnection class instance
    :param action: string giving the action type for the WebService
    :param kwargs: params to concatenate

    :rtype epured and updated dict
    """
    if isinstance(shibaconnection, ShibaConnection) is False:
        raise ShibaCallingError(
            "Internal parameter error : shibaconnection parameter is not "
            "a ShibaConnection instance"
        )
    if action not in shibaconnection.actionsinfo:
        raise ShibaCallingError(
            "Internal parameter error : action parameter "
            + action
            + " is unknown from the actions list"
        )
    newkwargs = {}
    for each in kwargs:
        if kwargs[each] is not None and kwargs[each] != "":
            newkwargs[each] = kwargs[each]
    newkwargs.update(shibaconnection.actionsinfo[action])
    newkwargs["action"] = action
    return newkwargs

def url_constructor(shibaconnection, params, domain=None):
    """URL constructor, formatting input and adding as many URL arguments as given.

    :param params: info dict, from inf_constructor
    :param domain: force URL domain, useful for some actions needing to go with http instead of https

    :rtype: formatted URL string
    """
    params.pop("self", None)
    category = params.pop("cat")
    ordered_params = OrderedDict(((k, params[k]) for k in sorted([k for k in params])))

    f = furl(domain or shibaconnection.domain)
    f.path.segments.append(category)
    f.path.normalize()
    f.add(ordered_params)
    return f.url

def _check_errors(obj):
    """Errors checking from the returned XML as object. Handling errors as internal exceptions raising

    :param obj: objectified element to check errors from

    :rtype: objectified element if errors have been found, False if the WebService hasn't encountered an error
    """
    if "errorresponse" in obj.tag:
        if "ServerError" == obj.error.code:
            raise ShibaParameterError(
                "Parameter error : "
                + obj.error.message
                + " - Reason : "
                + obj.error.details.detail
            )
        if "ParameterError" == obj.error.code:
            raise ShibaParameterError(
                "Parameter error : "
                + obj.error.message
                + " - Reason : "
                + obj.error.details.detail
            )
        if "InvalidUserConnection" == obj.error.code:
            raise ShibaLoginError(
                "Invalid user connection : "
                + obj.error.message
                + " - Reason : "
                + obj.error.details.detail
            )
        if "InvalidUserRights" == obj.error.code:
            if "Quota exceeded" in obj.error.message.text:
                raise ShibaQuotaExceededError(
                    "Too many requests : "
                    + obj.error.message
                    + " - Reason : "
                    + obj.error.details.detail
                )
            else:
                raise ShibaRightsError(
                    "Invalid user rights : "
                    + obj.error.message
                    + " - Reason : "
                    + obj.error.details.detail
                )
        return obj
    return False


#!/usr/bin/env python

""" Provision AWS Lex resources using python SDK
"""

# import time
# import boto3
from botocore.exceptions import ClientError

from intent_builder import IntentBuilder
# from slot_builder import SlotBuilder
from lex_helper import LexHelper
# from models.intent import Intent

class LexBotBuilder(LexHelper):

    MAX_DELETE_TRIES = 5
    RETRY_SLEEP = 5
    LOCALE = 'en-US'

    """Create/Update different elements that make up a Lex bot"""
    def __init__(self, logger, context, lex_sdk=None, intent_builder=None):
        self._logger = logger
        self._context = context
        if lex_sdk is None:
            self._lex_sdk = self._get_lex_sdk()
        else:
            self._lex_sdk = lex_sdk
        if intent_builder is None:

            self._intent_builder = IntentBuilder(self._logger, self._context, lex_sdk=self._lex_sdk)
        else:
            self._intent_builder = intent_builder

    def _replace_intent_version(self, bot_definition, intents):
        for intent in bot_definition['intents']:
            intent['intentVersion'] = intents[intent['intentName']]
        return bot_definition

    def _bot_put_properties(self, bot_name, messages, **kwargs):
        properties = {
            "name": bot_name,
            "locale": kwargs['locale'],
            "abortStatement": {
                "messages": [
                    {
                        "content": messages['abortStatement'],
                        "contentType": "PlainText"
                    }
                ]
            },
            "processBehavior": "BUILD",
            "childDirected": False,
            "clarificationPrompt": {
                "maxAttempts": 1,
                "messages": [
                    {
                        "content": messages['clarification'],
                        "contentType": "PlainText"
                    }
                ]
            },
            "description": kwargs['description'],
            "idleSessionTTLInSeconds": 3000
        }

        return properties

    def put(self, bot):
        """create bot"""
        intent_versions = self._put_intents(bot.name, bot.intents)
        self._logger.info(intent_versions)

        bot_response = self._put_bot(bot, intent_versions)
        return bot_response

    def delete(self, bot):
        """delete bot"""
        # TODO what about deleting published version(s) of the bot?
        self._delete_bot(bot.name)
        self._delete_intents(bot.name, bot.intents)

        self._logger.info('Successfully deleted bot and associated resources')

    def _put_intents(self, bot_name, intents):
        intent_versions = []
        for intent in intents:
            intent_versions.append(
                self._intent_builder.put_intent(intent)
            )

        return intent_versions

    def _delete_intents(self, bot_name, intents):
        intent_names = [intent.intent_name for intent in intents]
        # todo fix this so it just passes the intent object

        self._logger.info(intent_names)
        self._intent_builder.delete_intents(intent_names)

    def _bot_exists(self, name, versionOrAlias='$LATEST'):
        return self._get_resource(self._lex_sdk.get_bot,
                                  'get_bot',
                                  {'name': name, 'versionOrAlias': versionOrAlias})

    def _create_bot(self, bot_name, bot_properties):
        bot_exists, checksum = self._bot_exists(bot_name)
        if bot_exists:
            creation_response = self._update_lex_resource(
                self._lex_sdk.put_bot, 'put_bot', checksum, bot_properties)
            return creation_response, creation_response['checksum']

        else:
            self._logger.info(bot_properties)
            creation_response = self._create_lex_resource(
                self._lex_sdk.put_bot, 'put_bot', bot_properties)

            return creation_response, creation_response['checksum']

    def _put_bot(self, bot, intent_versions):
        """Create/Update bot"""
        self._logger.info('Put bot intent_versions %s', bot.name)
        self._logger.info(intent_versions)

        bot_properties = self._bot_put_properties(bot.name, bot.messages, **bot.attrs)
        bot_properties.update({"intents": intent_versions})

        self._logger.info("Bot properties for AWS %s", bot_properties)

        _, checksum = self._create_bot(bot.name, bot_properties)

        version_response = self._create_lex_resource(
            self._lex_sdk.create_bot_version, 'create_bot_version',
            {
                'name': bot.name,
                'checksum': checksum
            })

        self._logger.info("Created bot version %s", bot.name)
        self._logger.info(version_response)

        return version_response

    def _delete_bot(self, bot_name):
        '''Delete bot'''
        self._logger.info('deleting bot: %s', bot_name)
        while True:
            try:
                bot_exists, _ = self._bot_exists(bot_name)
                if(bot_exists):
                    self._delete_lex_resource(self._lex_sdk.delete_bot, 'delete_bot',
                                              name=bot_name)

                    self._logger.info('deleted bot: %s', bot_name)
                    break
                else:
                    break

            except ClientError:
                self._logger.warning('Lex can not call delete_bot on deleted bot %s.',
                                     bot_name)


from sklearn.model_selection import train_test_split as sklearn_train_test_split
from sklearn.preprocessing import LabelEncoder
from scipy.sparse import csr_matrix

import numpy as np
import pandas as pd

def get_examples(dataframe, labels_column="rating"):
    examples = dataframe[['userid', 'itemid']].values
    labels = dataframe[f'{labels_column}'].values
    return examples, labels

def train_test_split(examples, labels, test_size=0.1, verbose=0):
    if verbose:
        print("Train/Test split ")
        print(100-test_size*100, "% of training data")
        print(test_size*100, "% of testing data")

    train_examples, test_examples, train_labels, test_labels = sklearn_train_test_split(
        examples,
        labels,
        test_size=0.1,
        random_state=42,
        shuffle=True
    )

    train_users = train_examples[:, 0]
    test_users = test_examples[:, 0]

    train_items = train_examples[:, 1]
    test_items = test_examples[:, 1]

    x_train = np.array(list(zip(train_users, train_items)))
    x_test = np.array(list(zip(test_users, test_items)))

    y_train = train_labels
    y_test = test_labels

    if verbose:
        print()
        print('number of training examples : ', x_train.shape)
        print('number of training labels : ', y_train.shape)
        print('number of test examples : ', x_test.shape)
        print('number of test labels : ', y_test.shape)

    return (x_train, x_test), (y_train, y_test)

def ids_encoder(ratings):
    users = sorted(ratings['userid'].unique())
    items = sorted(ratings['itemid'].unique())

    uencoder = LabelEncoder()
    iencoder = LabelEncoder()

    uencoder.fit(users)
    iencoder.fit(items)

    ratings.userid = uencoder.transform(ratings.userid.tolist())
    ratings.itemid = iencoder.transform(ratings.itemid.tolist())

    return ratings, uencoder, iencoder

class ExplainableMatrixFactorization:

    def __init__(self, m, n, W, alpha=0.001, beta=0.01, lamb=0.1, k=10):
        self.W = W
        self.m = m
        self.n = n

        np.random.seed(64)

        self.k = k
        self.P = np.random.normal(size=(self.m,k))
        self.Q = np.random.normal(size=(self.n,k))

        self.alpha = alpha
        self.beta = beta
        self.lamb = lamb

        self.history = {
            "epochs":[],
            "loss":[],
            "val_loss":[],
        }

    def print_training_parameters(self):
        print('Training EMF')
        print(f'k={self.k} \t alpha={self.alpha} \t beta={self.beta} \t lambda={self.lamb}')

    def update_rule(self, u, i, error):
        self.P[u] = self.P[u] + \
            self.alpha*(2 * error*self.Q[i] - self.beta*self.P[u] - self.lamb*(self.P[u] - self.Q[i]) * self.W[u,i])

        self.Q[i] = self.Q[i] + \
            self.alpha*(2 * error*self.P[u] - self.beta*self.Q[i] + self.lamb*(self.P[u] - self.Q[i]) * self.W[u,i])

    def mae(self,  x_train, y_train):
        M = x_train.shape[0]
        error = 0
        for pair, r in zip(x_train, y_train):
            u, i = pair
            error += np.absolute(r - np.dot(self.P[u], self.Q[i]))
        return error/M

    def print_training_progress(self, epoch, epochs, error, val_error, steps=5):
        if epoch == 1 or epoch % steps == 0 :
                print(f"epoch {epoch}/{epochs} - loss : {round(error,3)} - val_loss : {round(val_error,3)}")

    def learning_rate_schedule(self, epoch, target_epochs = 20):
        if (epoch >= target_epochs) and (epoch % target_epochs == 0):
                factor = epoch // target_epochs
                self.alpha = self.alpha * (1 / (factor * 20))
                print("\nLearning Rate : {}\n".format(self.alpha))

    def fit(self, x_train, y_train, validation_data, epochs=10):
        self.print_training_parameters()

        x_test, y_test = validation_data

        for epoch in range(1, epochs+1):
            for pair, r in zip(x_train, y_train):
                u,i = pair
                r_hat = np.dot(self.P[u], self.Q[i])
                e = r - r_hat
                self.update_rule(u, i, error=e)

            error = self.mae(x_train, y_train)
            val_error = self.mae(x_test, y_test)
            self.update_history(epoch, error, val_error)
            self.print_training_progress(epoch, epochs, error, val_error, steps=1)

        return self.history

    def update_history(self, epoch, error, val_error):
        self.history['epochs'].append(epoch)
        self.history['loss'].append(error)
        self.history['val_loss'].append(val_error)

    def evaluate(self, x_test, y_test):
        error = self.mae(x_test, y_test)
        print(f"validation error : {round(error,3)}")

    def predict(self, userid, itemid, uencoder, iencoder):
        u = uencoder.transform([userid])[0]
        i = iencoder.transform([itemid])[0]
        r = np.dot(self.P[u], self.Q[i])
        return r

    def recommend(self, userid, uencoder, iencoder, N):
        u = uencoder.transform([userid])[0]

        predictions = np.dot(self.P[u], self.Q.T)

        top_idx = np.flip(np.argsort(predictions))[:N]

        top_items = iencoder.inverse_transform(top_idx)

        preds = predictions[top_idx]

        return top_items, preds

    def calc_mnap(self, k, df, users,uencoder, iencoder):
        if not 'itemid' in df.columns:
            errstr = "Column with items id's must be named itemid "
            raise TypeError(errstr)
        else:
            maps = []
            for user in users:
                score = 0.0
                num_hits = 0.0
                top_items, preds = self.recommend(user, uencoder, iencoder,N = 25)
                rel_items = df[df.userid == user]['itemid'].tolist()
                for i,p in enumerate(top_items[:k]):
                    if p in rel_items and p not in top_items[:i]:
                        num_hits += 1.0
                        score += num_hits / (i+1.0)
                maps.append(score / min(len(rel_items), k))
            return np.mean(maps)

    def calc_recalls(self, k, df, users, uencoder, iencoder):
        if not 'itemid' in df.columns:
            errstr = "Column with items id's must be named itemid "
            raise TypeError(errstr)
        else:
            recals = []
            for user in users:
                top_items, preds = self.recommend(user,uencoder, iencoder, N = 25)
                rel_items = df[df.userid == user]['itemid'].tolist()
                top_inter = 0
                for el in top_items[:5]:
                    if el in rel_items:
                        top_inter += 1
                recals.append(top_inter / len(top_items))
            return np.mean(recals)

    def dcg_at_k(self, r, k):
        r = r[:k]
        dcg = np.sum(r / np.log2(np.arange(2, len(r) + 2)))
        return dcg

    def ndgc_at_k(self, r, k, method = 0):
        dcg_max = self.dcg_at_k(sorted(r, reverse = True), k)
        return self.dcg_at_k(r, k) / dcg_max

    def calc_ndcg(self, predictions, k, users, ratings, movie_ids):
        ndgcs = []

        for target_users in np.unique(users):

            target_movie_id = movie_ids[target_users == users]
            target_rating = ratings[target_users == users]

            rel = target_rating[np.argsort(-predictions[target_users == users])]
            ndgc = self.ndgc_at_k(rel, k=k)
            ndgcs.append(ndgc)

        ndcg = np.mean(ndgcs)
        return ndcg


from osgeo import gdal
from osgeo import ogr
import numpy as np
import matplotlib.pyplot as plt
import os
import copy
import math
import argparse
from cal import calTiff
import config_this

##need to set to nodata after do atmospheric correction
def run(folder):

	#dump metadata ls8 to dictionary
	def metadata_ls8 (metadata):
		metadata = open(metadata)
		metadata_dict = dict()
		for line in metadata:
			if '    ' in line.split('=')[0]:
				try:
					D = {(line.split('=')[0]).replace('    ','').replace(' ',''):line.split('=')[1].replace(' ','').replace('\n','')}
					metadata_dict.update(D)
				except:
					'last line have no split'
		return metadata_dict

	#convert to reflectance
	def  reflectance (band,Mp,Ap,sinSE):
		band = gdal.Open(band)
		band.GetRasterBand(1)
		band = band.ReadAsArray().astype(np.float)
		band[np.where(band==0)] = np.nan
		band = (band*Mp+Ap)/sinSE
		return band
	#find F5bins as first bin that have frequency > 50 and next right to it also > 50
	def find_1st_valid_detail (band):
		bandVals = band.ravel()
		unique, counts = np.unique(bandVals, return_counts = True)

		count_next = np.append([0],counts[0:len(counts)-1])
		count_next_2 = np.append([0],count_next[0:len(count_next)-1])
		#idea is take (x1 - x) if >50 then (x2 - x1) >50 choose x1
		A = np.where(count_next-counts > 50)
		B = np.where(count_next_2-count_next > 50)
		C = np.array(B)[0] - np.array(A)[0]
		for i in C:
			if i == 1:
				break
		loc = np.array(A)[0][i]
		return unique[loc]
	def setNotadata_0(band, no):
	    band_buffer = band
	    outband = band.replace('1.TIF','2.TIF')
	    driver = gdal.GetDriverByName("GTiff")
	    band = gdal.Open(band)
	    band.GetRasterBand(1)
	    band = band.ReadAsArray().astype(np.float)
	    band[np.where(band<0)] = no
	    config_this.createTiff(band, band_buffer, outband)

	def do_ats_correct(folder, band):
		path_band = path['%s'%(band)]
		Fbin5_band =  find_1st_valid_detail(reflectance(path_band,Mp,Ap,sinSE))
		#((DN*Mp)+ Ap)/sinSE
		expression = '((%s*%f)+%f)/%f-%f' % (band, Mp, Ap, sinSE, Fbin5_band)
		out_filename = folder+'/'+os.path.basename(folder)+'_' + band + 'product1.TIF'
		calTiff (folder, expression, out_filename)
		setNotadata_0(out_filename,np.nan)

	#config section
	driver = gdal.GetDriverByName("GTiff")
	metadata = folder + '/' +os.path.basename(folder) + '_MTL.txt'
	metadata = metadata_ls8(metadata)
	SUN_ELEVATION = metadata['SUN_ELEVATION']
	SUN_ELEVATION = math.radians(float(SUN_ELEVATION))
	sinSE = math.sin(SUN_ELEVATION)

	#from landsat 8 metadata
	Mp = 0.00002
	Ap = -0.1

	Fbin5s = dict()
	path = dict()
	path['B2'] = folder + '/' +os.path.basename(folder) + '_B2.TIF'
	path['B3'] = folder + '/' +os.path.basename(folder) + '_B3.TIF'
	path['B4'] = folder + '/' +os.path.basename(folder) + '_B4.TIF'
	path['B5'] = folder + '/' +os.path.basename(folder) + '_B5.TIF'

	bands = ['B2','B3','B4','B5']

	for band in bands:
		do_ats_correct(folder,band)

def main ():
    parser = argparse.ArgumentParser()
    parser.add_argument("-f",dest = "folder",help = "insert folder need to do correction", type = str)
    args = parser.parse_args()
    run(args.folder)
   # run(arg.shp,args.date,args.province)
if __name__=='__main__':
    main()

#'LC08_L1TP_125052_20180217_20180307_01_T1'


import pandas as pd
from sklearn.linear_model import LinearRegression

df = pd.read_csv('heights.csv')
x = df['height']
y = df['weight']

line_fitter = LinearRegression()
data = x.values.reshape(-1, 1)
line_fitter.fit(data, y)

print(line_fitter.predict([[70]]))
print(line_fitter.coef_)
print(line_fitter.intercept_)

import matplotlib.pyplot as plt

plt.plot(x, y, 'o')
plt.plot(x, line_fitter.predict(data))
plt.show()


from wolframalpha import Client
#pip install wolframalpha

def setupWolfram():
    app_id='RWAGV2-6QGHW7Y95J'
    global client
    client=Client(app_id)

def startWolfram(question):
    try:
        res=client.query(question)
        res=next(res.results).text
        print(res)
        return res
    except:
        print("Not found on Wolfram alpha")
        return None

if __name__ == "__main__":
    setupWolfram()
    startWolfram("What is 2+3/4")


import os
from feature_extraction.date_utils import date_features
from feature_extraction.coords_features import coord_features
from feature_extraction.other_features import raw_features, categorical_features
from feature_extraction.path_utils import project_root

from metaflow import FlowSpec, step
import pandas as pd

class UbaarFeaturesExtraction(FlowSpec):
    @step
    def start(self):

        self.data = pd.read_csv(os.path.join(project_root(), 'data', 'raw', 'ubaar-competition', 'train.csv'),
                                encoding="utf-8", index_col="ID")
        self.features = pd.DataFrame()

        self.next(self.get_day_feature)

    @step
    def get_day_feature(self):
        self.features = date_features(self.data, self.features)

        self.next(self.get_lat_long_features)

    @step
    def get_lat_long_features(self):
        self.features = coord_features(self.data, self.features)

        self.next(self.add_categorical_features)

    @step
    def add_categorical_features(self):
        self.features = categorical_features(self.data, self.features)

        self.next(self.add_raw_features)

    @step
    def add_raw_features(self):

        self.features = raw_features(self.data, self.features)

        self.next(self.remove_missing_data)

    @step
    def remove_missing_data(self):

        self.features.dropna(inplace=True)
        self.next(self.report)

    @step
    def report(self):
        print(f"Shape: {self.features.shape}")
        self.next(self.save)

    @step
    def save(self):
        self.features.to_csv(os.path.join(project_root(), 'data', 'processed', 'ubaar_features.csv'))
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
    UbaarFeaturesExtraction()


import pathlib
import setuptools

here = pathlib.Path(__file__).parent.resolve()
readme = (here / 'README.md').read_text(encoding='utf-8')

# did not include torch and pyscf here
install_requires=['numpy', 'paramiko', 'ruamel.yaml', 'pytest-shutil', 'pydflow']

setuptools.setup(
    name="deepks2",
    version="1.2",
    setup_requires=['setuptools_scm'],
    author="Yifan Shan",
    author_email="xiaoshan@mail.ustc.edu.cn",
    description="DeePKS-flow: a 'dflow' workflow for DeePKS-kit",
    long_description=readme,
    packages=[
        "deepks2",
        "deepks2/flow",
        "deepks2/op/iter_op",
        "deepks2/step/iter_step",
        "deepks2/utils",
    ],
    long_description_content_type="text/markdown",
    # packages=setuptools.find_packages(include=['deepks2', 'deepks2.*']),
    classifiers=[
        "Programming Language :: Python :: 3.7",
    ],
    keywords='deepks2 DeePKS-flow',
    install_requires=install_requires,
    python_requires=">=3.7",
    entry_points={
        'console_scripts': [
            'deepks2=deepks2.main:main_cli',
            'dks2=deepks2.main:main_cli',
        ],
    },
)