import logging
import re
import xml.etree.cElementTree
import xml.sax.saxutils
from io import BytesIO
from typing import List, Optional, Sequence
from xml.etree.ElementTree import ElementTree

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader

logger = logging.getLogger(__name__)

_MAXIMUM_TITLE_LENGTH = 64

class QuipLoader(BaseLoader):
    """Load `Quip` pages.

    Port of https://github.com/quip/quip-api/tree/master/samples/baqup
    """

    def __init__(
        self, api_url: str, access_token: str, request_timeout: Optional[int] = 60
    ):
        """
        Args:
            api_url: https://platform.quip.com
            access_token: token of access quip API. Please refer:
            https://quip.com/dev/automation/documentation/current#section/Authentication/Get-Access-to-Quip's-APIs
            request_timeout: timeout of request, default 60s.
        """
        try:
            from quip_api.quip import QuipClient
        except ImportError:
            raise ImportError(
                "`quip_api` package not found, please run " "`pip install quip_api`"
            )

        self.quip_client = QuipClient(
            access_token=access_token, base_url=api_url, request_timeout=request_timeout
        )

    def load(
        self,
        folder_ids: Optional[List[str]] = None,
        thread_ids: Optional[List[str]] = None,
        max_docs: Optional[int] = 1000,
        include_all_folders: bool = False,
        include_comments: bool = False,
        include_images: bool = False,
    ) -> List[Document]:
        """
        Args:
            :param folder_ids: List of specific folder IDs to load, defaults to None
            :param thread_ids: List of specific thread IDs to load, defaults to None
            :param max_docs: Maximum number of docs to retrieve in total, defaults 1000
            :param include_all_folders: Include all folders that your access_token
                   can access, but doesn't include your private folder
            :param include_comments: Include comments, defaults to False
            :param include_images: Include images, defaults to False
        """
        if not folder_ids and not thread_ids and not include_all_folders:
            raise ValueError(
                "Must specify at least one among `folder_ids`, `thread_ids` "
                "or set `include_all`_folders as True"
            )

        thread_ids = thread_ids or []

        if folder_ids:
            for folder_id in folder_ids:
                self.get_thread_ids_by_folder_id(folder_id, 0, thread_ids)

        if include_all_folders:
            user = self.quip_client.get_authenticated_user()
            if "group_folder_ids" in user:
                self.get_thread_ids_by_folder_id(
                    user["group_folder_ids"], 0, thread_ids
                )
            if "shared_folder_ids" in user:
                self.get_thread_ids_by_folder_id(
                    user["shared_folder_ids"], 0, thread_ids
                )

        thread_ids = list(set(thread_ids[:max_docs]))
        return self.process_threads(thread_ids, include_images, include_comments)

    def get_thread_ids_by_folder_id(
        self, folder_id: str, depth: int, thread_ids: List[str]
    ) -> None:
        """Get thread ids by folder id and update in thread_ids"""
        from quip_api.quip import HTTPError, QuipError

        try:
            folder = self.quip_client.get_folder(folder_id)
        except QuipError as e:
            if e.code == 403:
                logging.warning(
                    f"depth {depth}, Skipped over restricted folder {folder_id}, {e}"
                )
            else:
                logging.warning(
                    f"depth {depth}, Skipped over folder {folder_id} "
                    f"due to unknown error {e.code}"
                )
            return
        except HTTPError as e:
            logging.warning(
                f"depth {depth}, Skipped over folder {folder_id} "
                f"due to HTTP error {e.code}"
            )
            return

        title = folder["folder"].get("title", "Folder %s" % folder_id)

        logging.info(f"depth {depth}, Processing folder {title}")
        for child in folder["children"]:
            if "folder_id" in child:
                self.get_thread_ids_by_folder_id(
                    child["folder_id"], depth + 1, thread_ids
                )
            elif "thread_id" in child:
                thread_ids.append(child["thread_id"])

    def process_threads(
        self, thread_ids: Sequence[str], include_images: bool, include_messages: bool
    ) -> List[Document]:
        """Process a list of thread into a list of documents."""
        docs = []
        for thread_id in thread_ids:
            doc = self.process_thread(thread_id, include_images, include_messages)
            if doc is not None:
                docs.append(doc)
        return docs

    def process_thread(
        self, thread_id: str, include_images: bool, include_messages: bool
    ) -> Optional[Document]:
        thread = self.quip_client.get_thread(thread_id)
        thread_id = thread["thread"]["id"]
        title = thread["thread"]["title"]
        link = thread["thread"]["link"]
        update_ts = thread["thread"]["updated_usec"]
        sanitized_title = QuipLoader._sanitize_title(title)

        logger.info(
            f"processing thread {thread_id} title {sanitized_title} "
            f"link {link} update_ts {update_ts}"
        )

        if "html" in thread:
            # Parse the document
            try:
                tree = self.quip_client.parse_document_html(thread["html"])
            except xml.etree.cElementTree.ParseError as e:
                logger.error(f"Error parsing thread {title} {thread_id}, skipping, {e}")
                return None

            metadata = {
                "title": sanitized_title,
                "update_ts": update_ts,
                "id": thread_id,
                "source": link,
            }

            # Download each image and replace with the new URL
            text = ""
            if include_images:
                text = self.process_thread_images(tree)

            if include_messages:
                text = text + "/n" + self.process_thread_messages(thread_id)

            return Document(
                page_content=thread["html"] + text,
                metadata=metadata,
            )
        return None

    def process_thread_images(self, tree: ElementTree) -> str:
        text = ""

        try:
            from PIL import Image
            from pytesseract import pytesseract
        except ImportError:
            raise ImportError(
                "`Pillow or pytesseract` package not found, "
                "please run "
                "`pip install Pillow` or `pip install pytesseract`"
            )

        for img in tree.iter("img"):
            src = img.get("src")
            if not src or not src.startswith("/blob"):
                continue
            _, _, thread_id, blob_id = src.split("/")
            blob_response = self.quip_client.get_blob(thread_id, blob_id)
            try:
                image = Image.open(BytesIO(blob_response.read()))
                text = text + "\n" + pytesseract.image_to_string(image)
            except OSError as e:
                logger.error(f"failed to convert image to text, {e}")
                raise e
        return text

    def process_thread_messages(self, thread_id: str) -> str:
        max_created_usec = None
        messages = []
        while True:
            chunk = self.quip_client.get_messages(
                thread_id, max_created_usec=max_created_usec, count=100
            )
            messages.extend(chunk)
            if chunk:
                max_created_usec = chunk[-1]["created_usec"] - 1
            else:
                break
        messages.reverse()

        texts = [message["text"] for message in messages]

        return "\n".join(texts)

    @staticmethod
    def _sanitize_title(title: str) -> str:
        sanitized_title = re.sub(r"\s", " ", title)
        sanitized_title = re.sub(r"(?u)[^- \w.]", "", sanitized_title)
        if len(sanitized_title) > _MAXIMUM_TITLE_LENGTH:
            sanitized_title = sanitized_title[:_MAXIMUM_TITLE_LENGTH]
        return sanitized_title

{'QuipClient': 'quip_api.quip.QuipClient', 'HTTPError': 'quip_api.quip.HTTPError', 'QuipError': 'quip_api.quip.QuipError'}._sanitize_title


import torch
from torch.autograd import Variable
from tensorboardX import SummaryWriter
from apex import amp
from utils import save_checkpoint, use_optimizer
from metrics import MetronAtK
import time

class Engine(object):
    """Meta Engine for training & evaluating NCF model

    Note: Subclass should implement self.model !
    """

    def __init__(self, config):
        self.config = config  # model configuration
        self._metron = MetronAtK(top_k=10)
        self._writer = SummaryWriter(log_dir='runs/{}'.format(config['alias']))  # tensorboard writer
        self._writer.add_text('config', str(config), 0)
        self.opt = use_optimizer(self.model, config)
        self.model, self.opt = amp.initialize(self.model,self.opt,opt_level='O2',loss_scale=32.0,combine_grad=True)
        self.crit = torch.nn.BCELoss()

    def train_single_batch(self, users, items, ratings):
        assert hasattr(self, 'model'), 'Please specify the exact model !'
        if self.config['use_cuda'] is True:
            users, items, ratings = users.npu(), items.npu(), ratings.npu()
        self.opt.zero_grad()
        ratings_pred = self.model(users, items)
        loss = self.crit(ratings_pred.view(-1), ratings)
        with amp.scale_loss(loss, self.opt) as scaled_loss:
            scaled_loss.backward()
        #loss.backward()
        self.opt.step()
        loss = loss.item()
        return loss

    def train_an_epoch(self, train_loader, epoch_id):
        assert hasattr(self, 'model'), 'Please specify the exact model !'
        self.model.train()
        time1 = time.time()
        total_loss = 0
        for batch_id, batch in enumerate(train_loader):
            assert isinstance(batch[0], torch.LongTensor)
            user, item, rating = batch[0], batch[1], batch[2]
            rating = rating.float()
            loss = self.train_single_batch(user, item, rating)
            print('[Training Epoch {}] Batch {}, Loss {}, sec/step {}'.format(epoch_id, batch_id, loss, time.time()-time1))
            total_loss += loss
        self._writer.add_scalar('model/loss', total_loss, epoch_id)

    def evaluate(self, evaluate_data, epoch_id):
        assert hasattr(self, 'model'), 'Please specify the exact model !'
        self.model.eval()
        with torch.no_grad():
            test_users, test_items = evaluate_data[0], evaluate_data[1]
            negative_users, negative_items = evaluate_data[2], evaluate_data[3]
            if self.config['use_cuda'] is True:
                test_users = test_users.npu()
                test_items = test_items.npu()
                negative_users = negative_users.npu()
                negative_items = negative_items.npu()
            test_scores = self.model(test_users, test_items)
            negative_scores = self.model(negative_users, negative_items)
            self._metron.subjects = [test_users.data.view(-1).tolist(),
                                 test_items.data.view(-1).tolist(),
                                 test_scores.data.view(-1).tolist(),
                                 negative_users.data.view(-1).tolist(),
                                 negative_items.data.view(-1).tolist(),
                                 negative_scores.data.view(-1).tolist()]
        hit_ratio, ndcg = self._metron.cal_hit_ratio(), self._metron.cal_ndcg()
        self._writer.add_scalar('performance/HR', hit_ratio, epoch_id)
        self._writer.add_scalar('performance/NDCG', ndcg, epoch_id)
        print('[Evluating Epoch {}] HR = {:.4f}, NDCG = {:.4f}'.format(epoch_id, hit_ratio, ndcg))
        return hit_ratio, ndcg

    def save(self, alias, epoch_id, hit_ratio, ndcg):
        assert hasattr(self, 'model'), 'Please specify the exact model !'
        model_dir = self.config['model_dir'].format(alias, epoch_id, hit_ratio, ndcg)
        save_checkpoint(self.model, model_dir)


import xarray as xr

import seaborn as sns

import warnings

def rates_to_df(arr):
    """
    Convert DataArray of rates to DataFrame

    Parameters
    ----------
    arr : xr.DataArray
        DataArray containing the rates of neurons across time

    Returns
    -------
    pd.DataFrame with neurons as index and time points as columns
    """
    non_time_dims = [dim for dim in arr.dims if dim != 'time']
    assert len(non_time_dims) == 1
    non_time_dim = non_time_dims[0]

    return (arr.to_dataframe('name')
               .reset_index()
               .pivot(index = non_time_dim, columns = 'time', values = 'name'))

def plot_rasterlike_rates(ax, rates, normalize=True, cbar_ax=None, cmap=None):
    """
    Plot rates of a trial similarly to a rasterplot

    Parameters
    ----------
    ax : plt.Axes
        Axes to plot on
    rates : pd.DataFrame
        rates to plot
        neuron x time
        convert DataArray to DataFrame by rates_to_df
    normalize : bool (default True)
        normalize every neuron's max. firing rate to 1
    cbar_ax : plt.Axes (default None)
        Axes to plot the colorbar on
        if None, don't plot a colorbar
    cmap : string or colormap
        colormap to use for the rates
    """
    if normalize:
        rates /= rates.max('time')

    if cbar_ax is not None:
        cbar = True
    else:
        cbar = False

    if cmap is None:
        cmap = sns.color_palette('Wistia', desat = 0.6)

    if isinstance(rates, xr.DataArray):
        plot_rates = rates_to_df(rates)
    else:
        plot_rates = rates
        warning.warn('would be nice to get an indexed DataArray')

    sns.heatmap(plot_rates, ax = ax, cbar = cbar, cbar_ax = cbar_ax,
                rasterized = True,
                cbar_kws = {'label' : 'rates'},
                cmap = cmap)


import base64
import gzip
import os
import re

import fsspec
import numpy as np
import pandas as pd
import pyarrow.parquet as pq

CLOUD_HEADER = 'gs://'

def do_external_setup(project_id):
  """Sets up GCP credentials for external colab.

  Args:
    project_id: GCP project id associated with the colab.
  """
  if not os.path.exists('/var/colab/hostname'):
    return
  with open('/var/colab/hostname') as f:
    vm_name = f.read()
  if vm_name.startswith('colab-'):
    del os.environ['NO_GCE_CHECK']
    del os.environ['GCE_METADATA_TIMEOUT']
  else:
    from google.colab import auth
    auth.authenticate_user()
  os.environ['GOOGLE_CLOUD_PROJECT'] = project_id

def format_plate_strings(plate_names):
  """Format the plate strings as strings of five digit ints.

  Args:
    plate_names: A pd.Series of strings representing the plate names that we
      want to format.

  Raises:
    ValueError: If plate_names contains a name that is more than five digits
      long.
  Returns:
    formatted_plates: A pd.Series representing the formatted plate names.
  """
  # Format the plate strings as 5 character strings with no decimal
  formatted_plates = plate_names.astype(str).apply(
      lambda x: x.split('.')[0].zfill(5))
  # If any of the plates are more than 5 digits, scream loudly.
  len_plate_names = np.array([len(p) for p in formatted_plates.values])
  if np.any(len_plate_names > 5):
    raise ValueError('Plate name > 5 characters found')
  # If any of the plates have non-digit characters, scream loudly.
  if not np.all([re.match(r'^\d+$', p) for p in formatted_plates.values]):
    raise ValueError('Plate with non-digit characters found')
  return formatted_plates

def read_csv_from_cloud(filepath):
  """Loads a csv into a dataframe from cloud storage.

  Any spaces or dashes in column names are converted to underscores so
  they play nice with pandas.

  Any 'plate' columns will be adjusted to be the 5 character string used
  by Google. Other columns will NOT be adjusted, you may need to do
  additional cleaning up because CSV does not save type information.

  Args:
    filepath: The cloud path to the file, starting from the bucket. Example -
      'bucket_name/folder_name/test.csv'

  Returns:
    pandas dataframe with the contents of the CSV.
  """
  cloud_filepath = os.path.join(CLOUD_HEADER, filepath)
  df = pd.read_csv(cloud_filepath)

  if 'plate' in df.columns:
    df['plate'] = format_plate_strings(df['plate'])

  column_names = list(df.columns)
  column_names = [x.replace(' ', '_') for x in column_names]
  column_names = [x.replace('-', '_') for x in column_names]
  df.columns = column_names

  return df

def to_csv_on_cloud(df, filepath):
  """Writes the dataframe to cloud as a CSV.

  Args:
    df: A pandas dataframe.
    filepath: A cloud bucket path, eg: 'bucket_name/folder_name/test.csv'
  """
  cloud_filepath = os.path.join(CLOUD_HEADER, filepath)
  with fsspec.open(cloud_filepath, 'w') as outfile:
    df.to_csv(outfile)

def to_file_on_cloud(txt, filepath):
  """Writes the given text to cloud.

  Args:
    txt: The text to write to the cloud.
    filepath: A cloud bucket path, eg: 'bucket_name/folder_name/test.html'
  """
  cloud_filepath = os.path.join(CLOUD_HEADER, filepath)
  with fsspec.open(cloud_filepath, 'w') as f:
    f.write(txt)

def read_file_df_from_cloud(filepath, dtype):
  """Reads dataframes saved as HDF5 or parquet format from cloud.

  Args:
    filepath: A cloud bucket path, eg: 'bucket_name/folder_name/test.html'
    dtype: data type the dataframe is saved as. Supports 'h5' and 'parquet'

  Returns:
    pd.DataFrame read from cloud
  """
  cloud_filepath = os.path.join(CLOUD_HEADER, filepath)

  # We can't read files directly from cloud;
  # I'll download it to a temp folder and load it here
  # fsspec.open_files handles glob if '*' is in filepath
  openfiles = fsspec.open_files(cloud_filepath)
  # Make a random string to attach to the filename to prevent people
  # from writing to the file at the same time
  randstr = str(np.random.random())[2:]
  df_shards = [pd.DataFrame()]
  for idx, openfile in enumerate(openfiles):
    with openfile as f:
      data = f.read()
    temp_path = 'temp_file_{}.file-{:05d}'.format(randstr, idx)
    with open(temp_path, 'wb') as fw:
      fw.write(data)
    if dtype == 'h5':
      df_shards.append(pd.read_hdf(temp_path))
    elif dtype == 'parquet':
      df_shards.append(pq.read_table(temp_path).to_pandas())
    os.remove(temp_path)
  return pd.concat(df_shards).sort_index()

def merge_dfs(df_list, keys_for_merge):
  """Joins all the dataframes on the keys.

  Does an outer join, does not check that all rows exist in all dataframes.

  Args:
    df_list: A list of pandas dataframes. Each needs to have the keys_for_merge
      as columns.
    keys_for_merge: A string list of column names to merge on. Typically
      ['plate', 'well'] aka WELL_MERGE_KEYS.

  Returns:
    The joined pandas dataframe.
  """
  if len(df_list) < 1:
    raise ValueError('I need at least 1 df to merge')

  for k in keys_for_merge:
    if k not in df_list[0]:
      raise ValueError('df missing column %s. Has: %s' %
                       (k, df_list[0].columns))
  merged_df = df_list[0]

  for df in df_list[1:]:
    cols_to_add = keys_for_merge + [c for c in df.columns if c not in merged_df]
    for k in keys_for_merge:
      if k not in df:
        raise ValueError('df missing column %s. Has: %s' % (k, df[0].columns))
    merged_df = merged_df.merge(df[cols_to_add], on=keys_for_merge, how='outer')

  return merged_df

def get_html_list(df_for_html, base_neuroglancer_url):
  """Creates the html for one pandas dataframe as a neuroglancer table.

  Args:
    df_for_html: The pandas dataframe to create the html from.
    base_neuroglancer_url: The string url for neuroglancer to load in the top
      iframe.

  Returns:
    A list of HTML strings that can be written to make the HTML page.
  """

  df_html = df_for_html.to_html(escape=False, table_id='link_df')
  encoded_html = base64.encodebytes(gzip.compress(
      df_html.encode('utf-8'))).decode().rstrip()

  include_html = """
  <link rel="stylesheet" type="text/css" href="https://cdn.datatables.net/1.11.3/css/jquery.dataTables.css">

  """
  # note backticks! these capture the \n's introduced by base64.encodebytes
  encoding_invocation = f'encoded_html=`{encoded_html}`;\n'
  invocation_js = encoding_invocation + """
  function stringify(x) {
    accum = [];
    for (var i=0; i<x.length; i+= 10000) {
        accum.push(String.fromCharCode.apply(null, x.slice(i, Math.min(i + 10000, x.length))));
    }
    return(accum.join(''));
  }
  function decode_base64_encoded_gzip(encoded_html) {
      gzip_encoded_bytes = new Uint8Array(atob(encoded_html).split('').map(function(x) {return x.charCodeAt(0)}));
      decoded_bytes = pako.inflate(gzip_encoded_bytes);
      decoded_bytes16 = new Uint16Array(decoded_bytes);
      decoded_html_str = stringify(decoded_bytes16);
      return decoded_html_str;
  }
  decoded_html_str = decode_base64_encoded_gzip(encoded_html);

  $('body').append($(decoded_html_str));
  $(document).ready( function () {
      $('#link_df').DataTable();
      $('#wait').hide();
  } );
  """

  accum = []
  accum.append(include_html)
  accum.append(
      f' '
  )
  accum.append('
')
  base_neuroglancer_link = ('Open '
                            'Neuroglancer') % (
                                base_neuroglancer_url)
  accum.append(base_neuroglancer_link)
  accum.append(
      '

Please wait while the table fully loads.

')
  accum.append(f'')
  return accum

def set_up_for_html(df, url_col_name):
  """Moves the url_col_name so it is the first column in the html output."""
  if url_col_name not in df.columns:
    raise ValueError('expected url col "%s" but did not find it.' %
                     (url_col_name))
  cols_without_url = list(df.columns)
  cols_without_url.remove(url_col_name)
  cols_in_order = [url_col_name] + cols_without_url
  df_for_html = df[cols_in_order]
  return df_for_html

def write_html_for_df(df,
                      out_path,
                      base_neuroglancer_url,
                      per_plate_pattern=None,
                      cloud_root_dir=''):
  """Writes a dataframe out as a neuroglancer table.

  Args:
    df: Pandas dataframe.
    out_path: String filepath for the output. Should be a cloud bucket path,
      eg -
        'path/to/folder/my_table.html'
    base_neuroglancer_url: The string url for neuroglancer to load in the top
      iframe.
    per_plate_pattern: Optional string filepath with one '%s'. If provided, the
      df must have a column "plate". One HTML file will be created for each
      plate.
    cloud_root_dir: Optional string filepath to folder containing the plate
      files. Use CLOUD_ROOT_DIR in colab.
  """
  accum = get_html_list(df, base_neuroglancer_url)
  to_file_on_cloud(''.join(accum), out_path)

  if per_plate_pattern:
    plates = list(df.plate.unique())
    for plate in plates:
      plate_df = df.query('plate == "%s"' % plate)
      plate_accum = get_html_list(plate_df, base_neuroglancer_url)
      plate_path = os.path.join(cloud_root_dir, per_plate_pattern % plate)
      to_file_on_cloud(''.join(plate_accum), plate_path)


# import mysql.connector
# from sqlalchemy import create_engine
import pandas as pd
import os
import numpy as np
import hashlib
import time
import humanize
import matplotlib.pyplot as plt
import seaborn as sns
import logging
logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
sns.set_theme(rc={'figure.dpi': 150, 'figure.figsize': (5, 3.75)})

import sys
# sys.path.append('/mnt/c/Users/akpau/Repos/data_proc')
sys.path.append('/home/apaulson/repos/data-proc/')

from mysql_cnx import create_mysql_cnx

###############################################################################
def get_hash(path, size):
    """ A helper function to generate the md5 checksum of a file.
    Args:
        path (str): path to file
        size (int): chunk size of file to read at a time
    Returns:
        hex digest of md5 sum.
    """
    m = hashlib.md5()
    with open(path, 'rb') as f:
        b = f.read(size)
        while len(b) > 0:
            m.update(b)
            b = f.read(size)
    return m.hexdigest()

###############################################################################
def draw_plate_heatmap(barcode='OTPARP1234-Eco1', df=None, wells='Well', values='SMDC_ID', labels='SMDC_ID', center=None, vmin=None, vmax=None, width=25, outdir=None):
    """
    Draw a heatmap of the plate using a platemap excel style file as your df
    Args:
        df (Pandas DataFrame): data to pivot into heatmap; optional if barcode is provided.
        row_col (str): name of column containing row values
        col_col (str): name of column containing column values
        values (str): name of column containing heatmap values (floats)
        labels (str): name of column containing heatmap label values (strings)
        barcode (str): name of plate or barcode to look up
    Returns:
        None
    Side effects:
        Draws a plate heatmap with values as labels
    """
    if df is None:
        df=make_platemaps_from_barcode(barcode=barcode, outdir=None)

    df['Row']=df[wells].str[0]
    df['Col']=df[wells].str[1:]
    df.Col=df.Col.astype(int)
    cbar=False
    pl  =  df.pivot_table(columns='Col', index='Row', values=values, dropna=False)
    pl  =  pl.astype(float)
    df[labels]=df[labels].astype(str).replace('nan','')
    plabels=df.pivot_table(columns='Col', index='Row', values=labels, aggfunc=lambda x: '\n'.join(x))
    for col in plabels.columns:
        plabels[col]=plabels[col].str.replace('.0', '', regex=False)
    if labels=='SMDC_ID':
        plabels=pl.copy()
        plabels=plabels.fillna(-99).astype(int).astype(str).replace('-99', '', regex=True)
        cbar=False

    # fig, ax = plt.subplots(1, figsize=(width,width/5))
    if cbar:
        fig, (ax, cax) = plt.subplots(1,2, figsize=(width,width/3),  gridspec_kw={"width_ratios":[1, .01]})
    else:
        fig, ax = plt.subplots(1, figsize=(width,width/3))
        cax=None
    sns.heatmap(pl, annot=plabels, fmt='s', linewidth=1, linecolor='white', ax=ax, cbar_ax = cax, cbar=cbar, center=center, vmin=vmin, vmax=vmax,
                annot_kws={"size": 7}
               )
    ax.set_title(f"{barcode} {values}");
    plt.tight_layout()
    if outdir is not None:
        plt.savefig(os.path.join(outdir, barcode+'.png'), bbox_inches='tight')
        plt.close()

###############################################################################
def make_platemaps_from_barcode(barcode='OTPARP1234-Eco1', outdir='./', cnx=None):
    """
    Create an excel file with each sheet showing a plate map with vendor alias,
    smdc_id, and lot for downstream reference.
    Args:
        barcode (str): Unique plate barcode to pull from database
        outdir (path): Place to save excel file. If None, no file will be saved.
    Returns:
        df: Long form table of all relevant plate information.
    """
    # pull data from hits
    close=False
    if cnx is None:
        close=True
        cnx = create_mysql_cnx()

    query=f"""
    select IW_COORDINATES, ICL_ALIAS as Vendor_Alias, ICL_IC_ID as SMDC_ID, IA_VALUE as Common_Name, IA_IAT_ID as Alias_Type,
        ICL_LOT_NUM as Lot_Num, IW_CONC_UM as Conc_uM, IW_FINAL_CONC_NM as Conc_nM, IC_SMILES as SMILES
    from INV_WELL left join INV_COMP_LOT on IW_ICL_ID = ICL_ID left join INV_ALIAS on IA_IC_ID=ICL_IC_ID left join INV_COMPOUND on IC_COMPOUND_ID=ICL_IC_ID
    where IW_IP_ID =
        (select IP_ID from INV_PLATE where IP_BARCODE = '{barcode}');
    """
    # with cnx.connect() as conn:
    df=pd.read_sql(query, cnx)
    if close:
        cnx.dispose()

    # filter for common name or keep nan
    wells=[]
    for well in df.IW_COORDINATES.unique():
        tmp=df[df.IW_COORDINATES==well]
        tmp=tmp[tmp.Alias_Type=='Common_Name'] # this will be one row
        if len(tmp)==0:
            tmp=pd.DataFrame(df[df.IW_COORDINATES==well].iloc[0]).T # this will be one row
            tmp['Common_Name']=np.nan
        wells.append(tmp)
    df=pd.concat(wells)
    try:
        df=df.drop(columns='Alias_Type')
    except:
        pass

    # format data
    df['Row']=df.IW_COORDINATES.str.slice(start=0, stop=1)
    df['Col']=df.IW_COORDINATES.str.slice(start=1).astype(int)
    df=df.rename(columns={'IW_COORDINATES':'Well'})
    df['Plate']=barcode
    cols = ['Vendor_Alias', 'SMDC_ID', 'Common_Name', 'Lot_Num']

    # return or save
    if outdir is None:
        return df

    # save as plate layouts
    else:
        with pd.ExcelWriter(os.path.join(outdir, f'{barcode}_plate_maps.xlsx')) as writer:
            for col in cols:
                if col in ['Vendor_Alias','Common_Name']:
                    df[col]=df[col].fillna('None')
                    aggfunc=lambda x: ' '.join(x)
                else:
                    df[col]=df[col].fillna(-99).astype(int)
                    aggfunc='mean'
                tmpdf=df.pivot_table(index='Row', columns='Col', values=col, aggfunc=aggfunc)
                tmpdf=tmpdf.replace('None', np.nan).replace(-99, np.nan)
                tmpdf.index.name=col
                tmpdf.to_excel(writer, sheet_name=col)
            df=df.replace('None', np.nan).replace(-99, np.nan)
            df.to_excel(writer, sheet_name='Long_Form')
        return df

###############################################################################
def lookup_lot_by_user(ids=['285699','130958'] ,users=[("Amanda","Paulson")]):
    """ Find the most recent (highest) lotnum that the user has created for their compound.
    Args:
        ids (list of strings): The list of SMDC IDs to look up. Integers, but in string form.
        users (list of string tuples): The list of users in the form ("FirstName","LastName")
    Returns:
        fullresults (DataFrame): the SMDC IDs and Lots as two columns in a df.
    """
    # pull data from hits
    cnx = create_mysql_cnx()
    res=[]
    for fname, lname in users:
        print('looking up users..')
        query=f"""select DM_ID from DIRECTORY_MEMBER where DM_FNAME='{fname}' and DM_LNAME='{lname}';"""
        # with cnx.connect() as conn:
        user=pd.read_sql(query, cnx)
        user=user.DM_ID.iloc[0]

        print('looking up compounds..')
        format_strings = ','.join(['%s'] * len(ids))

        query=f"""
        select IC_COMPOUND_ID as SMDC_ID, max(ICL_LOT_NUM) as Lot from INV_COMPOUND
        inner join INV_COMP_LOT on IC_COMPOUND_ID=ICL_IC_ID
        where IC_COMPOUND_ID in (%s) and ICL_INS_BY={user}
        group by IC_COMPOUND_ID;
        """ % format_strings
        fullresults=pd.read_sql(query, cnx, params=tuple(ids))
        res.append(fullresults)
    cnx.dispose()

    # create df
    print('creating df..')
    fullresults=pd.concat(res).reset_index(drop=True)
    return fullresults

###############################################################################
def lookup_lot_by_lot_num(ids=['285699','130958'] ,lot_nums=['1','1']):
    """ Find the most recent (highest) lotnum that the user has created for their compound.
    Args:
        ids (list of strings): The list of SMDC IDs to look up. Integers, but in string form.
        lot_nums (list of integers): The list of lots to look up. Integers, but in string form.
    Returns:
        fullresults (DataFrame): the SMDC IDs and Lots as two columns in a df.
    """
    # pull data from hits
    cnx = create_mysql_cnx()
    res=[]
    for fname, lname in users:
        print('looking up users..')
        query=f"""select DM_ID from DIRECTORY_MEMBER where DM_FNAME='{fname}' and DM_LNAME='{lname}';"""
        # with cnx.connect() as conn:
        user=pd.read_sql(query, cnx)
        user=user.DM_ID.iloc[0]

        print('looking up compounds..')
        format_strings = ','.join(['%s'] * len(ids))

        query=f"""
        select IC_COMPOUND_ID as SMDC_ID, max(ICL_LOT_NUM) as Lot from INV_COMPOUND
        inner join INV_COMP_LOT on IC_COMPOUND_ID=ICL_IC_ID
        where IC_COMPOUND_ID in (%s) and ICL_INS_BY={user}
        group by IC_COMPOUND_ID;
        """ % format_strings
        fullresults=pd.read_sql(query, cnx, params=tuple(ids))
        res.append(fullresults)
    cnx.dispose()

    # create df
    print('creating df..')
    fullresults=pd.concat(res).reset_index(drop=True)
    return fullresults

###############################################################################
# def lookup_plates_by_ico_id(id=135):
#     cnx = create_mysql_cnx()
#     query=f"""
#     select IP_ID from INV_PLATE where IP_ICO_ID in

###############################################################################
def read_platemap_block(plate_map=None, block=0, value_name='SMDC_ID'):
    """ Read a block formatted 384w plate map from excel.
    Args:
        plate_map (DataFrame): the dataframe containing the block formatted data
        block (int): the numerical ID of the block to read, 0 at the top
    Returns:
        df (DataFrame): the new dataframe containing the long-format block data.
        df columns are row, col and value_name
    """
    cols = ['row']
    cols.extend(range(1,25))
    df = plate_map.iloc[17*block:17*block+16, 0:25]
    df.columns=cols
    df = df.melt(id_vars = ['row'], value_vars=range(1,25), ignore_index=True,
                 value_name=value_name, var_name='col')
    df.col=df.col.astype(str).str.rjust(2, '0')
    return df

###############################################################################
def make_dr_platemaps(dp, row, users, basedir):
    """ Construct dr platemap from info provided in excel file.
    TODO: refactor this so the text can be printed in top function... or add logging...
    """

    cnx = create_mysql_cnx()
    print(row.Grid_num)

    # destination plate info
    #smdc ids
    dids = read_platemap_block(plate_map=dp[row.Grid_num[0:31]], block=0, value_name='SMDC_ID')
    dids=dids.replace('DMSO', np.nan)
    if sum(dids.SMDC_ID.astype(str).str.contains('-'))>0:
        dids[['SMDC_ID','Lot']]=dids.SMDC_ID.astype(str).str.split('-', expand=True)
        dids.SMDC_ID=dids.SMDC_ID.replace('nan', np.nan).replace(u'\xa0', np.nan)
        dids.SMDC_ID=dids.SMDC_ID.astype(float)
        dids.Lot=dids.Lot.astype(float)

    dids['Destination well']=dids.row+dids.col
    dids['Destination plate'] = row.Grid_num
    if row.Raw_assay_ID != 'multiple':
        dids['Assay ID'] = int(row.Raw_assay_ID)
    else:
        aids=read_platemap_block(plate_map=dp[row.Grid_num[0:31]], block=2, value_name='Assay ID')
        aids['Destination well']=aids.row+aids.col
        dids=dids.merge(aids[['Destination well','Assay ID']], how='left')
    if ('Lot_info' in row.index):
        if (row.Lot_info == 'use map'):
            aids=read_platemap_block(plate_map=dp[row.Grid_num[0:31]], block=2, value_name='Lot')
            aids['Destination well']=aids.row+aids.col
            dids=dids.merge(aids[['Destination well','Lot']], how='left')

    dums= read_platemap_block(plate_map=dp[row.Grid_num[0:31]], block=1, value_name='conc')
    dums['Destination well']=dums.row+dums.col
    dids=dids.merge(dums[['Destination well','conc']], how='left')

    # source plate info
    if isinstance(row.Source_plate, str):
        if 'map' in row.Source_plate:
            sps=read_platemap_block(plate_map=dp[row.Grid_num[0:31]], block=2, value_name='Source plate')
            sps['Destination well']=sps.row+sps.col
            dids=dids.merge(sps[['Destination well','Source plate']])
            source_plates=sps[~sps['Source plate'].isna()]['Source plate'].unique().tolist()
        elif 'lot' in row.Source_plate:
            aids=read_platemap_block(plate_map=dp[row.Grid_num[0:31]], block=2, value_name='Lot')
            aids['Destination well']=aids.row+aids.col
            dids=dids.merge(aids[['Destination well','Lot']], how='left')
            source_plates=[]
        else:
            source_plates = row.Source_plate.split(';')
            source_plates = [x.strip() for x in source_plates]
            source_plates = [x for x in source_plates if x!='']
        # else:
        #     source_plates =
    if len(source_plates)>0:
        dfs=[]
        for plate in source_plates:
            df=make_platemaps_from_barcode(barcode=plate, outdir=None, cnx=cnx)
            df=df[~df.SMDC_ID.isna()]
            dfs.append(df)
        df=pd.concat(dfs)
        df.columns=['Source well', 'alias', 'SMDC_ID', 'Common Name','Lot', 'Conc_uM', 'Conc_nM', 'SMILES','Row', 'Col', 'Source plate']
        dids = dids.merge(df[['SMDC_ID', 'Lot', 'Source plate', 'Source well']], how='left')
    else:
        dids['Source plate']='NA'
        dids['Source well']='NA'

    if 'Lot' not in dids.columns:
            logging.warning("no Lot from plate map")
            dids['Lot']=np.nan
    if len(dids[(~dids.SMDC_ID.isna())&(dids.Lot.isna())])>0:
        ids=dids[~dids.SMDC_ID.isna()].SMDC_ID.astype(int).astype(str).unique().tolist()
        lots=lookup_lot_by_user(ids=ids, users=users)
        dids.loc[(~dids.SMDC_ID.isna())&(dids.Lot.isna()),'Lot']=dids.loc[(~dids.SMDC_ID.isna())&(dids.Lot.isna()),'SMDC_ID'].map(dict(zip(lots.SMDC_ID,lots.Lot)))
        dids['Source plate']='NA'
        dids['Source well']='NA'

    dids=dids.fillna('NA')

    # check for zero concentrations
    if 0 in dids.conc.unique():
        minconc=input(f"There are zero's in the concentration column! {dids.conc.sort_values(ascending=False).round(4).unique().tolist()}.\nWhat would you like to replace 0 with? (Enter NA to ignore 0's): ")
        try:
            minconc=float(minconc)
            dids.loc[dids.conc==0, 'conc']=minconc
        except:
            dids=dids[dids.conc!=0]
        print(f"Min conc of plate is now {dids.conc.min()}")
    dids=dids.rename(columns={'conc':'[compound] uM'})
    dids=dids[['SMDC_ID', 'Lot', 'Source plate', 'Source well', 'Destination plate',
               'Destination well', '[compound] uM', 'Assay ID']]
    cnx.dispose()

    # integrity check
    if len(dids) not in [96, 384, 1536]:
        logging.warning(f"Warning: your plate has {len(dids)} wells. Was a compound pulled from two plates?")

    # fill empty with NA
    dids=dids.fillna('NA')
    return dids

###############################################################################
def make_dr_platemaps_and_folders(infile='', users=[("Amanda","Paulson")], basedir='./', dose_response=True):
    """
    Create an excel file with dose response plate information for DR assay upload
    to HiTS. Each sheet in the input file should be labeled with the destination
    plate barcode and have SMDC IDs and Concs. There should be a sheet labeled 'Info'
    with additional data including source plate barcodes, assay IDs and dates.
    Args:
        barcode (str): Unique plate barcode to pull from database
        basedir (path): Place to move data to and create folders in.
    Returns:
        df: Long form table of all relevant plate information.
    """

    # read plate layout file
    if 'xlsx' in infile:
        dp=pd.read_excel(infile, engine='openpyxl', sheet_name=None)
    else:
        dp={}
        dp['Info']=pd.read_csv(infile)

    info=dp['Info']
    info=info.dropna(how='all', axis='columns')
    info=info.groupby([x for x in info.columns if 'Grid_num' not in x]).agg(lambda x: ','.join(x)).reset_index()

    for assay, date in zip(info.Raw_assay_ID, info.Date_of_screen):
        df=info[(info.Raw_assay_ID==assay)&(info.Date_of_screen==date)].reset_index(drop=True)
        for i, row in df.iterrows():
            row.Date_of_screen=str(int(row.Date_of_screen))
            # create folders
            outdir=os.path.join(basedir, str(int(row.Raw_assay_ID))+'_'+row.Assay_name, f'toload_{i}', row.Date_of_screen)
            if not os.path.exists(outdir):
                os.makedirs(outdir)
                print(f'create folder {outdir}')
            # make and save platemaps if dose response
            if dose_response:
                dids=make_dr_platemaps(dp, row, users, basedir)
                dids.to_excel(os.path.join(outdir, f'{row.Grid_num}_DR_plate.xlsx'), index=False)

            # create or modify grid_num.txt
            gridnum=pd.DataFrame(row.Grid_num.split(','), columns=['grid_num'])
            if os.path.exists(os.path.join(outdir, 'grid_num.txt')):
                gn=pd.read_csv(os.path.join(outdir, 'grid_num.txt'))
                gn=pd.concat([gn,gridnum])
                gn=gn.drop_duplicates(keep='first')
            else:
                gn=gridnum
            gn.to_csv(os.path.join(outdir, 'grid_num.txt'), index=False)

            # move data files to folders
            print('Copying data...')
            starttime=time.time()
            row.Raw_data_file = row.Raw_data_file.replace('\\','/').replace('Z:','').replace('Y:','').replace('/SMDC','SMDC').replace('/shared','').replace('sharedSMDC','SMDC').replace('//montreal.ucsf.edu','').replace('smb:','')
            # origfile=os.path.join('/mnt/z', row.Raw_data_file)
            origfile=os.path.join('/mnt/mac/Volumes/Shared/', row.Raw_data_file)
            # print(origfile)
            fname=row.Raw_data_file.rsplit('/',1)[-1]
            # print(fname)
            # outfile=os.path.join(outdir.replace('/Shared','').replace('/shared',''), fname)
            outfile=os.path.join(outdir, fname)
            # print(outfile)
            # if os.path.exists(outfile):
            #     # cmd="test $(md5sum %s %s | awk '{print $1}' | uniq | wc -l) == 1" % (outfile, origfile)
            #     # res=os.system(cmd)
            #     # if res==1:
            #     # if get_hash(outfile, 65536) != get_hash(origfile, 65536):
            #     cmd=f"cp '{origfile}' '{outdir.replace('/Shared','')}/'"
            #     os.system(cmd)
            # else:
            # cmd=f"cp '{origfile}' '{outdir.replace('/Shared','').replace('/shared','')}/'"
            cmd=f"cp '{origfile}' '{outdir}/'"
            os.system(cmd)
            print(f'Done copying {humanize. precisedelta(starttime - time.time(), minimum_unit="seconds", format="%d")}')

# ###############################################################################
# inv_well_cols=['IW_SOURCE_ID','IW_IP_ID','IW_BARCODE','IW_WELL_INDEX','IW_COORDINATES','IW_ICL_ID','IW_VOL_UL','IW_CONC_UM','IW_FINAL_CONC_NM','IW_DILUTION_FACTOR','IW_FREEZETHAW','IW_INS_BY','IW_INS_DATE','IW_CONTAM','IW_NOTES']
# def make_well_upload_list(pm):

###############################################################################
def clean_up_folders(infile=None, basedir='./', delete_orig=False):
    """ After data is uploaded into HiTS, delete files from all_data (optional)
    and move each dated folder from toload to loaded.
    Args:
        infile (path): path to the plate map info file
        basedir (path): path to the directory for which to move data
        delete_orig (bool): whether to delete the original data from its location
    """
    # read data
    if 'xlsx' in infile:
        dp=pd.read_excel(infile, engine='openpyxl', sheet_name='Info')
    else:
        dp=pd.read_csv(infile)
    dp=dp[["Raw_assay_ID","Assay_name","Date_of_screen","Raw_data_file"]].drop_duplicates()
    # delete original file from all_data or similar folder
    if delete_orig:
        for rdf in dp.Raw_data_file:
            rdf = rdf.replace('\\','/').replace('Z:','').replace('Y:','').replace('/SMDC','SMDC').replace('/shared','').replace('sharedSMDC','SMDC').replace('//montreal.ucsf.edu','').replace('smb:','')
            rdf=os.path.join('/mnt/mac/Volumes/Shared/',rdf)
            print(rdf)
            if os.path.exists(rdf):
                cmd=f"rm {rdf}"
                os.system(cmd)

    # move each date from toload to loaded folders

    for assay, date in zip(dp.Raw_assay_ID, dp.Date_of_screen):
        df=dp[(dp.Raw_assay_ID==assay)&(dp.Date_of_screen==date)].reset_index(drop=True)
        for i, row in df.iterrows():
            outdir=os.path.join(basedir, str(row.Raw_assay_ID)+'_'+row.Assay_name)
            toload=os.path.join(outdir, f'toload_{i}')
            loaded=os.path.join(outdir, 'loaded')
            if not os.path.exists(loaded):
                os.mkdir(loaded)
            print(loaded)
            for datadir in os.listdir(toload):
                oldpath=os.path.join(toload, datadir)
                oldpath=oldpath.replace('(','\(').replace(')','\)').replace(' ','\ ')
                loaded = loaded.replace('(','\(').replace(')','\)').replace(' ','\ ')
                cmd=f"mv {oldpath} {loaded}/"
                os.system(cmd)
                cmd=f"rm -r {oldpath}"
                os.system(cmd)


# Databricks notebook source
from IPython.display import display
from datetime import datetime
import ipywidgets as widgets
from scenario_class import Scenario
from databricks.sdk.runtime import spark
import pandas as pd
import numpy as np

scenario = Scenario()

# COMMAND ----------

optimisation = scenario.optimisation.df
assumption = scenario.assumption.df.set_index('period_end').sort_index()

PC_VSC_flow = optimisation['PC_VSC_flow']
PC_solar_production = optimisation['PC_solar_production']
operational_capacity = assumption.loc[assumption['identifier'] == 'PC_battery_operational_capacity']

BESS_df = pd.DataFrame()
num_inverters = 1554

BESS_df['Demand/Inverter'] = PC_VSC_flow.div(num_inverters)
BESS_df['PV Production/Inverter'] = PC_solar_production.div(num_inverters)
BESS_df['BESS Size'] = 20

# bess_index = []
# oc_index = operational_capacity.index
# for index in BESS_df.index:
#     i = index.strftime("%Y-%m-%d")
#     if i in oc_index:
#         BESS_df['BESS Size'].loc[index] = operational_capacity['numeric_value'].loc[i]

BESS_df['BESS Size'] = BESS_df['BESS Size']
BESS_df['Net Flow of Energy'] = BESS_df['PV Production/Inverter'] - BESS_df['Demand/Inverter']

charge_conditions = [BESS_df['Net Flow of Energy'] >= 0, BESS_df['Net Flow of Energy'] < 0]
charge_choices = ['Charge', 'Discharge']
BESS_df['Charge/Discharge'] = np.select(charge_conditions, charge_choices)

bess_states = []
spill = []
bess_state = 20
bess_effiency = 0.95
inverter_effiency = 0.95
bess_power = 2
current_spill = 0
for row in BESS_df.iterrows():
    net_flow = row[1][3]
    if bess_state + net_flow >= 0 and bess_state + net_flow <= bess_size:
        if net_flow > bess_power:
            current_spill += net_flow - bess_power
            bess_state += bess_power*bess_effiency
        else:
            bess_state += net_flow*bess_effiency
        bess_states.append(bess_state)
    spill.append(current_spill)

BESS_df['BESS State of Charge'] = b_size
BESS_df['Spill'] = spill
BESS_df

# COMMAND ----------

BESS_df.tail(30)

# COMMAND ----------

BESS_df.plot(y='Spill', use_index=True)

# COMMAND ----------

solcast_query = "select * from sandbox.pv_data.solcast__18_201125_133_408077"
df_solcast = spark.sql(solcast_query).toPandas()
optimisation = scenario.optimisation.df

num_rows = len(optimisation.index)
df_solcast = df_solcast.iloc[:num_rows]
df_solcast = df_solcast.drop(columns=['PeriodEnd', 'PeriodStart'])
optimisation = optimisation.reset_index(drop=True)
df = optimisation.join(df_solcast, how='outer')
df


"""Module with multilayer perceptron definition.

"""

import json
import pathlib
import time
from copy import deepcopy
from typing import List, Optional

import numpy as np

from .activations import _ACTIVATION_DICT
from .metrics import _METRICS_DICT
from .utils import chunked, progress_bar, unison_shuffle

class MultilayerPerceptron:
    """Multilayer perceptron estimator.

    Args:
        inputs: Number of inputs to multilayer perceptron.
        units: List of numbers of units in consecutive layers of perceptron.
        activations: List of names of activation functions for consecutive
                     layers. If not given all layers will have sigmoid
                     activation functions.

    Attributes:
        _model: Model data for estimator.

    """

    def __init__(
            self, inputs: int, units: List[int], activations: Optional[List[str]] = None
        ) -> None:
        if activations is None:
            activations = ['sigmoid' for _ in range(len(units))]
        else:
            if len(units) != len(activations):
                raise ValueError('length of units and activations are different')
            for activation in activations:
                if activation not in _ACTIVATION_DICT:
                    raise ValueError(f'{activation} not available')

        layer_shapes = zip((inputs, *units)[1:], (inputs, *units)[:-1])
        self._activations: List[str] = activations
        self._layers: List[np.ndarray] = [
            np.random.rand(units_num, inputs_num + 1) * 2 - 1
            for units_num, inputs_num in layer_shapes]

    def __str__(self) -> str:
        """Creates model summary in form of string.

        Returns:
            Model summary in form of string.

        """
        text = f'{"_"*48}\n{"inputs":<16}{"units(outputs)":<16}{"activation":<16}\n{"="*48}'
        for layer, activation in zip(self._layers, self._activations):
            text += f'\n{layer.shape[1]-1:<16}{layer.shape[0]:<16}{activation:<16}\n{"_"*48}'
        return text

    def predict(self, x: np.ndarray) -> np.ndarray:
        """Predicts outputs for given inputs.

        Args:
            x: Inputs for prediction in form of numpy array of vectors.

        Returns:
            Prediction in form of numpy array of vectors.

        """
        self._validate_x(x)

        y = _bias_extend(np.array(x, ndmin=2))
        for layer, activation in zip(self._layers, self._activations):
            y = _bias_extend(_ACTIVATION_DICT[activation](np.matmul(layer, y.T).T))
        y = _bias_reduce(y)
        return y

    def train(
            self,
            train_x: np.ndarray,
            train_y: np.ndarray,
            val_x: Optional[np.ndarray] = None,
            val_y: Optional[np.ndarray] = None,
            batch_size: int = 32,
            epochs: int = 20,
            learning_rate: float = 0.01,
            momentum: float = 0.9,
            shuffle: bool = True,
            metrics: str = 'mse',
            early_stopping: bool = True,
            verbose: bool = True
        ) -> None:
        """Trains estimator on given data.

        Args:
            train_x: Training input data (2d).
            train_y: Training labels (2d).
            val_x: Validation input data (2d).
            val_y: Validation labels (2d).
            batch_size: Number of vectors in single batch.
            epochs: Number of training epochs.
            learning_rate: Learning rate for model.
            momentum: Momentum or model.
            shuffle: Flag for shuffling input data every epoch.
            metrics: Name of metrics to evaluate model with.
            early_stopping: Flag for early stopping.
            verbose: Flag for training verbosity.

        """
        # pylint: disable=too-many-arguments, too-many-locals, too-many-branches, too-many-statements
        # Check if metrics is available
        if metrics not in _METRICS_DICT:
            raise ValueError(f'{metrics} is not available')
        # Validate data
        self._validate_x(train_x)
        self._validate_y(train_y)
        if val_x is not None and val_y is not None:
            self._validate_x(val_x)
            self._validate_y(val_y)
        # Print training summary
        if verbose:
            print(f'Train on {len(train_x)} samples', end='')
            if val_x is not None and val_y is not None:
                print(f' - Validate on {len(val_x)} samples')
            else:
                print()
        # Calculate base scores
        train_score, compare_score = _METRICS_DICT[metrics](self.predict(train_x), train_y)
        if val_x is not None and val_y is not None:
            val_score, compare_score = _METRICS_DICT[metrics](self.predict(val_x), val_y)
        # Save best model and score
        if early_stopping:
            best_score = compare_score
            best_activations = deepcopy(self._activations)
            best_layers = deepcopy(self._layers)

        deltas = [np.zeros_like(layer) for layer in self._layers]
        for epoch_i in range(epochs):
            if verbose:
                start_time = time.time()
                print(f'Epoch {epoch_i+1}/{epochs}')
            if shuffle:
                train_x, train_y = unison_shuffle(train_x, train_y)

            for batch_x, batch_y in progress_bar(
                    zip(chunked(train_x, batch_size), chunked(train_y, batch_size)),
                    total=len(train_x),
                    step=batch_size,
                    verbose=verbose):
                # Feedforward and save results
                # `inputs` -> layer -> `outputs` -> activation
                inputs, outputs = [_bias_extend(batch_x)], []
                for i, (layer, activation) in enumerate(zip(self._layers, self._activations)):
                    outputs.append(np.matmul(layer, inputs[i].T).T)
                    inputs.append(_bias_extend(_ACTIVATION_DICT[activation](outputs[i])))
                batch_y_pred = _bias_reduce(inputs.pop())
                # Backpropagate
                errors = [
                    (batch_y_pred - batch_y)
                    * _ACTIVATION_DICT[self._activations[-1]](outputs.pop(), derivative=True)]
                for i, (layer, activation) in enumerate(zip(
                        reversed(self._layers), reversed(self._activations[:-1]))):
                    errors.append(
                        np.matmul(_bias_reduce(layer).T, errors[i].T).T
                        * _ACTIVATION_DICT[activation](outputs.pop(), derivative=True))
                errors = list(reversed(errors))
                # Update weigths
                deltas = [
                    -learning_rate * np.matmul(input_.T, error).T + momentum * delta
                    for input_, error, delta in zip(inputs, errors, deltas)]
                for layer, delta in zip(self._layers, deltas):
                    layer += delta

            # Calculate scores
            train_score, compare_score = _METRICS_DICT[metrics](self.predict(train_x), train_y)
            if val_x is not None and val_y is not None:
                val_score, compare_score = _METRICS_DICT[metrics](self.predict(val_x), val_y)
            # Save best model and score
            was_better = False
            if early_stopping and compare_score >= best_score:
                was_better = True
                best_score = compare_score
                best_activations = deepcopy(self._activations)
                best_layers = deepcopy(self._layers)
            # Print epoch summary
            if verbose:
                duration = time.time() - start_time
                # Print results
                if val_x is not None and val_y is not None:
                    print(
                        f'duration: {duration:.2f}s - train_{metrics}: {train_score:.4f} '
                        f'- val_{metrics}: {val_score:.4f}',
                        end='')
                else:
                    print(f'duration: {duration:.2f}s - train_{metrics}: {train_score:.4f}', end='')
                if early_stopping and was_better:
                    print(' <<')
                else:
                    print()

        # Restore best model
        if early_stopping:
            self._activations = best_activations
            self._layers = best_layers

    def save(self, file_path: pathlib.Path) -> None:
        """Saves model to file.

        Args:
            file_path: Path to file to which model will be saved.

        """
        model = {
            'activations': self._activations,
            'layers': [layer.tolist() for layer in self._layers]}
        with open(pathlib.Path(file_path), 'w', encoding='utf-8') as model_file:
            json.dump(model, model_file, indent=4, sort_keys=True)

    @staticmethod
    def load(file_path: pathlib.Path) -> 'MultilayerPerceptron':
        """Loads model from file.

        Args:
            file_path: Path to model file.

        """
        with open(pathlib.Path(file_path), 'r', encoding='utf-8') as model_file:
            model = json.loads(model_file.read())
            multilayer_perceptron = MultilayerPerceptron(inputs=1, units=[1,])
            multilayer_perceptron._activations = model['activations']
            multilayer_perceptron._layers = [np.array(layer, ndmin=2) for layer in model['layers']]
            return multilayer_perceptron

    def _validate_x(self, x: np.ndarray) -> None:
        """Validates x data regarding model.

        Args:
            x: Data to validate.

        Raises:
            ValueError: In case of failed validation.

        """
        if x.ndim != 2:
            raise ValueError('x should have 2 dims')
        if x.shape[1] != self._layers[0].shape[1] - 1:
            raise ValueError(
                f'invalid x shape (None, {x.shape[1]}), '
                f'should be (None, {self._layers[0].shape[1] - 1})')

    def _validate_y(self, y: np.ndarray) -> None:
        """Validates y data regarding model.

        Args:
            y: Data to validate.

        Raises:
            ValueError: In case of failed validation.

        """
        if y.ndim != 2:
            raise ValueError('y should have 2 dims')
        if y.shape[1] != self._layers[-1].shape[0]:
            raise ValueError(
                f'invalid y shape (None, {y.shape[1]}), '
                f'should be (None, {self._layers[-1].shape[0]})')

def _bias_extend(array: np.ndarray) -> np.ndarray:
    """Extends data with ones.

    Method extends each vector in array with additional one at the end.

    Args:
        array: Array to extend.

    Returns:
        Array extended with ones.

    """
    return np.hstack((array, np.ones((array.shape[0], 1))))

def _bias_reduce(array: np.ndarray) -> np.ndarray:
    """Reduces data.

    Method removes last element from each vector in array.

    Args:
        array: Array to reduce.

    Returns:
        Reduced array.

    """
    return array[:, :-1]


from django.test import TestCase
from .models import Watering
from datetime import timedelta
from django.utils import timezone
from .api import viewsets
from django.urls import reverse

class WateringLogic(TestCase):

    def test_check_soil_needs_watering_time(self):
        watering, next_interval = viewsets.check_last_watering_interval()
        print("next_interval", next_interval)
        self.assertEqual(watering, True)

    def test_post_dry_humidity_to_api(self):
        url = reverse('jardim_api:humidity')
        data = {'humidity': 900}
        response = self.client.post(url, data=data)
        watering = response.json()['watering']
        duration = response.json()['duration']
        print(response.status_code)
        print("watering", watering)
        print("duration", duration)
        self.assertEqual(response.status_code, 200)
        self.assertEqual(watering, True)
        self.assertGreater(duration, 0)

    def test_watering_under_interval(self):
        watering = Watering()
        watering.time_stamp = timezone.now()
        watering.save()
        url = reverse('jardim_api:humidity')
        data = {'humidity': 900}
        response = self.client.post(url, data=data)
        watering = response.json()['watering']
        duration = response.json()['duration']
        print(response.status_code)
        print("watering", watering)
        print("duration", duration)
        self.assertEqual(response.status_code, 200)
        self.assertEqual(watering, False)
        self.assertEqual(duration, 0)

    def test_watering_above_interval(self):
        watering = Watering()
        time = timezone.now() - timedelta(minutes=21)
        watering.time_stamp = time
        watering.save()
        url = reverse('jardim_api:humidity')
        data = {'humidity': 900}
        response = self.client.post(url, data=data)
        watering = response.json()['watering']
        duration = response.json()['duration']
        print(response.status_code)
        print("watering", watering)
        print("duration", duration)
        self.assertEqual(response.status_code, 200)
        self.assertEqual(watering, True)
        self.assertGreater(duration, 0)


import re
import datetime
import win10toast
import threading

def find_closest_match(pattern, text):
    # Find all matches of the pattern in the text
    matches = re.findall(pattern, text)

    # Get the current date and time
    now = datetime.datetime.now()

    # Set the minimum difference to the maximum possible value
    min_difference = datetime.timedelta.max

    # Set the closest match to None
    closest_match = None

    # Iterate over the matches
    for match in matches:
        # Parse the date and time from the match
        date = datetime.datetime.strptime(match, '%Y-%m-%d %H:%M:%S')

        # Calculate the difference between the current date and time and the match
        difference = abs(date - now)

        # If the difference is smaller than the minimum difference, set the closest match to the current match
        if difference < min_difference:
            min_difference = difference
            closest_match = match

    # Return the closest match
    return closest_match

def extract_command(file):
    # Read the Markdown file
    with open(file, 'r') as f:
        text = f.read()

    # Extract the command from the text
    pattern = r'\*\*(.+?)\*\*'
    match = find_closest_match(pattern, text)

    if match:
        # Split the command into time and message
        time, message = match.group(1).split(';')

        # Strip leading and trailing whitespace from the time and message
        time = time.strip()
        message = message.strip()

        return time, message

    return None, None

def launch_alarm(file):
    # Create a ToastNotifier object
    toaster = win10toast.ToastNotifier()

    # Extract the command from the Markdown file
    time, message = extract_command(file)

    if time and message:
        # Parse the time and display a notification
        alarm_time = datetime.datetime.strptime(time, '%Y-%m-%d %H:%M:%S')
        toaster.show_toast("Alarm", message, alarm_time)
    else:
        print("Command not found")

def run_alarm_thread(file):
    # Create a ToastNotifier object
    toaster = win10toast.ToastNotifier()

    # Extract the command from the Markdown file
    time, message = extract_command(file)

    # Set the current alarm to None
    current_alarm = None

    # Run the loop indefinitely
    while True:
        # If the current alarm is None or has already gone off, extract the next command from the file
        if not current_alarm or current_alarm < datetime.datetime.now():
            time, message = extract_command(file)
            current_alarm = datetime.datetime.strptime(time, '%Y-%m-%d %H:%M:%S')

        # Calculate the time until the next alarm
        time_until_alarm = current_alarm - datetime.datetime.now()

        # Sleep for the time until the next alarm
        time.sleep(time_until_alarm.total_seconds())

        # Display the notification
        toaster.show_toast("Alarm", message, current_alarm)
def run_alarm(file):
    # Create a ToastNotifier object
    toaster = win10toast.ToastNotifier()

    # Extract the command from the Markdown file
    time, message = extract_command(file)

    # Set the current alarm to None
    current_alarm = None

    # Run the loop indefinitely
    while True:
        # If the current alarm is None or has already gone off, extract the next command from the file
        if not current_alarm or current_alarm < datetime.datetime.now():
            time, message = extract_command(file)
            current_alarm = datetime.datetime.strptime(time, '%Y-%m-%d %H:%M:%S')

        # Calculate the time until the next alarm
        time_until_alarm = current_alarm - datetime.datetime.now()

        # Sleep for the time until the next alarm
        time.sleep(time_until_alarm.total_seconds())

        # Display the notification
        toaster.show_toast("Alarm", message, current_alarm)

# Run the alarm
#run_alarm('alarm.md')

# Create a thread to run the alarm in the background
#alarm_thread = threading.Thread(target=run_alarm_thread, args=('alarm.md',))

# Start the thread
#alarm_thread.start()

#launch_alarm('alarm.md')

# example of md file
#**2022-01-01 00:00:00; Wake up!**
#It's time to start your day.


#!/usr/bin/env python3
import socket
import numpy as np
import json
import matplotlib.pyplot as plt
import matplotlib.animation as animation

HOST = '<IP>'  # IP address
PORT = 65431  # Port to listen on (use ports > 1023)
DATA_SIZE = 50 # number of data to be displayed in the graph simultaneously, you can change it to any number you want

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.bind((HOST, PORT))
    s.listen()
    print('Server is listening...')
    conn, addr = s.accept()
    with conn:
        print('Connected by', addr)

        # set up the figure
        fig = plt.figure(figsize=plt.figaspect(0.5))

        # create 3D subplots for gyro and accelerometer data
        ax_gyro = fig.add_subplot(1, 2, 1, projection='3d')
        ax_acce = fig.add_subplot(1, 2, 2, projection='3d')

        # initialize empty lists for storing data
        gyro_x = [0]
        gyro_y = [0]
        gyro_z = [0]
        acce_x = [0]
        acce_y = [0]
        acce_z = [0]

        # define the animation function
        def animate(frame):
            message = "server requesting data"
            conn.sendall(message.encode('utf-8'))
            data = conn.recv(1024).decode('utf-8')
            print('Received from socket server : ', data)
            # print data
            json_data = json.loads(data)
            # json_formatted_str = json.dumps(json_data, indent=2)
            # print(json_formatted_str)
            # append data to list
            gyro_x.append(json_data["Gyro"]["x"])
            gyro_y.append(json_data["Gyro"]["y"])
            gyro_z.append(json_data["Gyro"]["z"])
            acce_x.append(json_data["Acce"]["x"])
            acce_y.append(json_data["Acce"]["y"])
            acce_z.append(json_data["Acce"]["z"])

            # Limit the data lists to a maximum length of 100
            if len(gyro_x) > DATA_SIZE:
                gyro_x.pop(0)
                gyro_y.pop(0)
                gyro_z.pop(0)
                acce_x.pop(0)
                acce_y.pop(0)
                acce_z.pop(0)
            # time.sleep(3)
            ax_gyro.clear()
            ax_gyro.plot(gyro_x, gyro_y, gyro_z, label="Gyro", color="blue")
            ax_gyro.legend()
            ax_acce.clear()
            ax_acce.plot(acce_x, acce_y, acce_z, label="Acce", color="red")
            ax_acce.legend()

        # create a FuncAnimation object and display the plot
        ani = animation.FuncAnimation(fig, animate, frames=range(0, DATA_SIZE), interval=200)
        plt.show()


from pathlib import Path

import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset

class PPFDataset(Dataset):
    def __init__(self, dataset_dir, num_images=None):
        assert Path(f'data/{dataset_dir}/labels.csv').is_file()
        df = pd.read_csv(f'data/{dataset_dir}/labels.csv')

        if num_images == None:
            num_images = len(df)
        else:
            assert num_images < len(df)
        num_patches = pd.Series(df['num_patches'][:num_images], dtype='int32')
        filename = pd.Series(df['filename'][:num_images])

        self.filenames = [
            f'{fname}_{i:03d}.npy'
            for fname, imax in zip(filename, num_patches)
            for i in range(imax)
        ]

        assert Path(f'data/{dataset_dir}/image').is_dir()
        self.ds_dir = f'data/{dataset_dir}'

        self.dataset_size = len(self.filenames)

    def __getitem__(self, idx):
        assert Path(f'{self.ds_dir}/image/{self.filenames[idx]}').is_file()

        with open(f'{self.ds_dir}/image/{self.filenames[idx]}', 'rb') as file:
            img_ppf = np.load(file)

        return img_ppf

    def __len__(self):
        return self.dataset_size

class PPFPairDataset(Dataset):
    def __init__(self, dataset_dir, num_images=None):
        assert Path(f'data/{dataset_dir}/labels.csv').is_file()
        df = pd.read_csv(f'data/{dataset_dir}/labels.csv')

        if num_images == None:
            num_images = len(df)
        else:
            assert num_images < len(df)
        num_patches = pd.Series(df['num_patches'][:num_images], dtype='int32')
        filename = pd.Series(df['filename'][:num_images])

        self.filenames = [
            f'{fname}_{i:03d}.npy'
            for fname, imax in zip(filename, num_patches)
            for i in range(imax)
        ]

        assert Path(f'data/{dataset_dir}/image').is_dir()
        assert Path(f'data/{dataset_dir}/model').is_dir()
        self.ds_dir = f'data/{dataset_dir}'

        self.dataset_size = len(self.filenames)

    def __getitem__(self, idx):
        assert Path(f'{self.ds_dir}/image/{self.filenames[idx]}').is_file()
        with open(f'{self.ds_dir}/image/{self.filenames[idx]}', 'rb') as file:
            img_ppf = np.load(file)

        assert Path(f'{self.ds_dir}/model/{self.filenames[idx]}').is_file()
        with open(f'{self.ds_dir}/model/{self.filenames[idx]}', 'rb') as file:
            md_ppf = np.load(file)

        return img_ppf, md_ppf

    def __len__(self):
        return self.dataset_size

if __name__ == '__main__':
    from statistics import mean, median, stdev
    from torch.utils.data import DataLoader
    from tqdm import tqdm
    from configs import DEVICE

    train_ds = PPFPairDataset('ycbv_obj_000001_train', 50)
    train_dl = DataLoader(train_ds, batch_size=1)
    print(f'Training set size: {train_ds.__len__()}')

    x, y = next(iter(train_dl))
    print(x[0, :5, :], y[0, :5, :], sep='\n')
    print(x.size(), y.size())

    loss_func = torch.nn.CosineSimilarity(dim=0)
    loss = loss_func(x, y)
    print(loss.size(), loss, sep='\n')

    # losses = [
    #     torch.sum(loss_func(x.to(DEVICE), y.to(DEVICE))).item()
    #     for x, y in tqdm(train_dl)
    # ]
    # print(f'Mean: {mean(losses):.3f} / Median: {median(losses):.3f} / Stdev: {stdev(losses):.3f}')


from algo.cnn import CNN

import os
from pathlib import Path
import torch
import torch.nn as nn
import torch.optim as optim
from torch import autograd
import numpy as np

class Discriminator(nn.Module):  # Discriminator Network

    # initializer
    def __init__(self, device, lr, state_dim, num_actions, wasserstein, grad_penalty, branched=False):
        super(Discriminator, self).__init__()
        self.lr = lr
        self.state_dim = state_dim
        self.num_actions = num_actions
        self.device = device
        self.branched = branched

        self.grad_penalty = grad_penalty
        self.wasserstein = wasserstein
        self.bce_loss = nn.BCEWithLogitsLoss()

        self.cnn = CNN(n_channels=9)
        if self.branched:
            self.disc = nn.Linear(512+1+num_actions, 3)
        else:
            self.disc = nn.Linear(512+3+1+num_actions, 1)

        self.optimizer = optim.AdamW(self.parameters(), lr=lr, weight_decay=0.1)
        self.to(self.device)

        checkpoint_dir = os.path.join(Path(__file__).parent.parent, 'checkpoints')
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        self.checkpoint_file = os.path.join(checkpoint_dir, f'bc_gail_disc')

    # forward propagation
    def forward(self, state, command, speed, action):
        cnn_out = self.cnn(state)
        if self.branched:
            disc_inp = torch.cat([cnn_out, speed, action], dim=1)
            preds = self.disc(disc_inp)
            pred = command[:, 0:1] * preds[:, 0:1] + \
                   command[:, 1:2] * preds[:, 1:2] + \
                   command[:, 2:3] * preds[:, 2:3]
        else:
            disc_inp = torch.cat([cnn_out, command, speed, action], dim=1)
            pred = self.disc(disc_inp)
        return pred

    # gradient penalty (only works with single-branched architecture)
    def compute_grad_pen(self, expert_states, expert_commands, expert_speeds, expert_actions,
                         agent_states, agent_commands, agent_speeds, agent_actions, lambda_=10):

        alpha = torch.rand(expert_states.size(0), 1)

        exp_cnn_out = self.cnn(expert_states)
        agt_cnn_out = self.cnn(agent_states)
        expert_data = torch.cat([exp_cnn_out, expert_commands, expert_speeds, expert_actions], dim=1)
        policy_data = torch.cat([agt_cnn_out, agent_commands, agent_speeds, agent_actions], dim=1)
        alpha = alpha.expand_as(expert_data).to(expert_data.device)
        mixup_data = alpha * expert_data + (1 - alpha) * policy_data
        disc = self.disc(mixup_data)

        ones = torch.ones(disc.size()).to(self.device)
        grad = autograd.grad(outputs=disc, inputs=mixup_data, grad_outputs=ones,
                             create_graph=True, retain_graph=True, only_inputs=True)[0]

        grad_pen = lambda_ * (grad.norm(2, dim=1) - 1).pow(2).mean()
        return grad_pen

    def learn(self, disc_batch_size, expert_states, expert_commands, expert_speeds, expert_actions,
              agent_states, agent_commands, agent_speeds, agent_actions):

        n_states = agent_states.shape[0]

        # sample expert data
        expert_states_sample, expert_commands_sample, expert_speeds_sample, expert_actions_sample = sample_expert(
            expert_states, expert_commands, expert_speeds, expert_actions, n_states)

        # creating tensors from expert states and actions
        expert_states_sample = torch.from_numpy(expert_states_sample)
        expert_commands_sample = torch.from_numpy(expert_commands_sample)
        expert_speeds_sample = torch.from_numpy(expert_speeds_sample)
        expert_actions_sample = torch.from_numpy(expert_actions_sample)

        # creating mini-batches
        batch_starts = np.arange(0, n_states, disc_batch_size)

        mean_real_loss = 0
        mean_fake_loss = 0
        mean_raw_loss = 0
        mean_loss = 0

        for b in batch_starts:
            expert_states_batch = expert_states_sample[b: b + disc_batch_size].to(self.device)
            expert_commands_batch = expert_commands_sample[b: b + disc_batch_size].to(self.device)
            expert_speeds_batch = expert_speeds_sample[b: b + disc_batch_size].to(self.device)
            expert_actions_batch = expert_actions_sample[b: b + disc_batch_size].to(self.device)
            agent_states_batch = agent_states[b: b + disc_batch_size].to(self.device)
            agent_commands_batch = agent_commands[b: b + disc_batch_size].to(self.device)
            agent_speeds_batch = agent_speeds[b: b + disc_batch_size].to(self.device)
            agent_actions_batch = agent_actions[b: b + disc_batch_size].to(self.device)

            # calculating discriminator prediction
            agent_preds = torch.squeeze(
                self.forward(agent_states_batch, agent_commands_batch,
                             agent_speeds_batch, agent_actions_batch), dim=-1)

            expert_preds = torch.squeeze(
                self.forward(expert_states_batch, expert_commands_batch,
                             expert_speeds_batch, expert_actions_batch), dim=-1)

            if self.wasserstein:
                raw_disc_loss = -(expert_preds.mean() - agent_preds.mean())
            else:

                # real and fake labels (1s and 0s)
                real_labels = torch.ones(expert_states_batch.shape[0]).to(self.device)
                fake_labels = torch.zeros(agent_states_batch.shape[0]).to(self.device)

                # calculating the loss
                loss_real = self.bce_loss(expert_preds, real_labels)
                loss_fake = self.bce_loss(agent_preds, fake_labels)
                raw_disc_loss = loss_real + loss_fake

            if self.grad_penalty:
                grad_pen = self.compute_grad_pen(
                    expert_states_batch, expert_commands_batch, expert_speeds_batch, expert_actions_batch,
                    agent_states_batch, agent_speeds_batch, agent_commands_batch, agent_actions_batch)
                disc_loss = raw_disc_loss + grad_pen
            else:
                disc_loss = raw_disc_loss

            # backpropagation and optimization
            self.optimizer.zero_grad()
            disc_loss.backward()
            self.optimizer.step()

            mean_real_loss += loss_real.item() if not self.wasserstein else 0.0
            mean_fake_loss += loss_fake.item() if not self.wasserstein else 0.0
            mean_raw_loss += raw_disc_loss.item()
            mean_loss += disc_loss.item()

        mean_real_loss /= len(batch_starts)
        mean_fake_loss /= len(batch_starts)
        mean_raw_loss /= len(batch_starts)
        mean_loss /= len(batch_starts)

        return mean_real_loss, mean_fake_loss, mean_raw_loss, mean_loss

    # generate rewards
    def generate_rewards(self, agent_states, agent_commands, agent_speeds, agent_actions):
        with torch.no_grad():
            probs = torch.sigmoid(self.forward(agent_states, agent_commands, agent_speeds, agent_actions))
            clamped_probs = torch.clamp(probs, min=0.01, max=0.99)
            rewards = - torch.log(1 - clamped_probs)
        return rewards.squeeze()

    def save_models(self):
        torch.save(self.state_dict(), self.checkpoint_file)

    def load_models(self):
        self.load_state_dict(torch.load(self.checkpoint_file, map_location=self.device))

def sample_expert(expert_states, expert_commands, expert_speeds, expert_actions, num):
    indices = np.random.choice(np.arange(expert_states.shape[0]), num, replace=False)
    return expert_states[indices], expert_commands[indices], expert_speeds, expert_actions[indices]


# from byol_essential.byol_models.model import BYOL
import os
import math
import numpy as np
import torch
import torchvision.transforms as transforms
from data.augment import Augment, Cutout
from utils.collate import collate_custom
import pdb
import copy

def get_criterion(p):
    if p['criterion'] == 'simclr':
        from losses.losses import SimCLRLoss
        criterion = SimCLRLoss(**p['criterion_kwargs'])

    elif p['criterion'] == 'scan':
        from losses.losses import SCANLoss
        criterion = SCANLoss(**p['criterion_kwargs'])

    elif p['criterion'] == 'confidence-cross-entropy':
        from losses.losses import ConfidenceBasedCE
        criterion = ConfidenceBasedCE(p['confidence_threshold'], p['criterion_kwargs']['apply_class_balancing'])

    elif p['criterion'] in  ['propmix']:
        from losses.losses import SemiLoss
        return SemiLoss()

    else:
        raise ValueError('Invalid criterion {}'.format(p['criterion']))

    return criterion

def get_feature_dimensions_backbone(p):
    if p['backbone'] == 'resnet18':
        return 512

    elif p['backbone'] == 'resnet50':
        return 2048

    elif p['backbone'] == 'InceptionResNetV2':
        return 1536

    else:
        raise NotImplementedError

def get_model(p, pretrain_path=None):
    # Get backbone
    if p['backbone'] == 'resnet18':
        if p['train_db_name'] in ['cifar-10', 'cifar-20', 'cifar-100', 'mini_imagenet_blue', 'mini_imagenet_red', 'mini_imagenet32_red','stanford_cars_blue']:
            from models.resnet_cifar import resnet18
            backbone = resnet18()

        # elif p['train_db_name'] == 'stl-10':
        #     from models.resnet_stl import resnet18
        #     backbone = resnet18()

        else:
            raise NotImplementedError
    elif p['backbone'] == 'PRN18':
        from models.preact_resnet_cifar import PreActResNet18
        backbone = PreActResNet18()
    elif p['backbone'] == 'vgg19':
        raise NotImplementedError

    elif p['backbone'] == 'densenet':
        from models.densenet_rog import densenet

        backbone = densenet(num_classes=p['num_classes'])

    elif p['backbone'] == 'resnet50':
        if 'imagenet' in p['train_db_name']:
            from models.resnet import resnet50
            backbone = resnet50()

        else:
            raise NotImplementedError

    elif p['backbone'] == 'InceptionResNetV2':
        if  p['train_db_name'] in ['webvision', 'imagenet_100']:
        # if  p['train_db_name'] in ["imagenet_100",]:
            from models.InceptionResNetV2 import network
            backbone = network()

        else:

            raise NotImplementedError
    elif p['backbone'] == 'resnet34':
        if p['train_db_name'] in ['cifar-10-plc', 'cifar-100-plc']:
            from models.resnet import resnet34
            backbone = resnet34()

    else:
        import pdb; pdb.set_trace()
        raise ValueError('Invalid backbone {}'.format(p['backbone']))

    # Setup
    if p['setup'] in ['simclr', 'moco']:
        from models.models import ContrastiveModel
        model = ContrastiveModel(backbone, **p['model_kwargs'])

    elif p['setup'] in ['scan']:
        from models.models import ClusteringModel
        # if p['setup'] == 'selflabel':
        #     assert(p['num_heads'] == 1)
        model = ClusteringModel(backbone, p['num_classes'], p['num_heads'], p['setup'])

    elif p['setup'] in ['propmix']:
        from models.models import Model
        model = Model(backbone, p['num_classes'], p['num_heads'], p['setup'])

    else:
        raise ValueError('Invalid setup {}'.format(p['setup']))

    # Load pretrained weights

    if pretrain_path is not None and os.path.exists(pretrain_path):
        state = torch.load(pretrain_path, map_location='cpu')

        if p['setup'] == 'scan': # Weights are supposed to be transfered from contrastive training
            missing = model.load_state_dict(state, strict=False)

        # elif p['setup'] == 'selflabel': # Weights are supposed to be transfered from scan
        #     # We only continue with the best head (pop all heads first, then copy back the best head)
        #     model_state = state['model']
        #     all_heads = [k for k in model_state.keys() if 'cluster_head' in k]
        #     best_head_weight = model_state['cluster_head.%d.weight' %(state['head'])]
        #     best_head_bias = model_state['cluster_head.%d.bias' %(state['head'])]
        #     for k in all_heads:
        #         model_state.pop(k)

        #     model_state['cluster_head.0.weight'] = best_head_weight
        #     model_state['cluster_head.0.bias'] = best_head_bias
        #     missing = model.load_state_dict(model_state, strict=True)

        elif p['setup'] in ['propmix']: # Weights are supposed to be transfered from scan
            # We only continue with the best head (pop all heads first, then copy back the best head)

            model_state = state['model']

            all_heads = [k for k in model_state.keys() if 'cluster_head' in k]
            best_head_weight = model_state['cluster_head.%d.weight' %(state['head'])]
            best_head_bias = model_state['cluster_head.%d.bias' %(state['head'])]
            for k in all_heads:
                model_state.pop(k)

            model_state['head.weight'] = best_head_weight
            model_state['head.bias'] = best_head_bias
            # model_state['sl_head.weight'] = best_head_weight
            # model_state['sl_head.bias'] = best_head_bias
            missing = model.load_state_dict(copy.deepcopy(model_state), strict=True)

        else:
            raise NotImplementedError

    elif pretrain_path is not None and not os.path.exists(pretrain_path):
        raise ValueError('Path with pre-trained weights does not exist {}'.format(pretrain_path))

    else:
        pass

    return model

def get_train_dataset(p, transform, to_augmented_dataset=False,
                        to_neighbors_dataset=False, to_noisy_dataset=False, split=None, meta_info={}):
    # Base dataset
    if p['train_db_name'] == 'cifar-10':
        if p['setup'] in ['propmix']:
            from data.cifar_propmix import cifar_dataset
            dataset = cifar_dataset(dataset=p['dataset'],root_dir=p['data_path'],transform=transform, meta_info=meta_info)
            to_noisy_dataset = False
        else:
            from data.cifar import CIFAR10
            dataset = CIFAR10(root=p['data_path'] , train=True, transform=transform, download=True)

    elif p['train_db_name'] == 'cifar-100':
        if p['setup'] in [ 'propmix']:
            from data.cifar_propmix import cifar_dataset
            dataset = cifar_dataset(dataset=p['dataset'],root_dir=p['data_path'],transform=transform, meta_info=meta_info)
            to_noisy_dataset = False
        else:
            from data.cifar import CIFAR100
            dataset = CIFAR100(root=p['data_path'], train=True, transform=transform, download=True)

    elif p['train_db_name'] == 'imagenet':
        from data.imagenet import ImageNet
        dataset = ImageNet(split='train', transform=transform)

    elif p['train_db_name'] in ['imagenet_50', 'imagenet_100', 'imagenet_200']:
        from data.imagenet import ImageNetSubset
        subset_file = './data/imagenet_subsets/%s.txt' %(p['train_db_name'])
        dataset = ImageNetSubset(subset_file=subset_file, split='train', transform=transform)

    elif p['train_db_name'] == 'webvision':
        from data.webvision import Webvision
        dataset = Webvision(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'])
    elif p['train_db_name'] == 'mini_imagenet_blue':
        from data.redblue import MiniImagenet
        dataset = MiniImagenet(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'], color='blue')
    elif p['train_db_name'] in  ['mini_imagenet_red', 'mini_imagenet32_red']:
        from data.redblue import MiniImagenet
        dataset = MiniImagenet(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'], color='red')
    elif p['train_db_name'] == 'stanford_cars_blue':
        from data.redblue import StanfordCars
        dataset = StanfordCars(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'], color='blue')
    elif p['train_db_name'] == 'cifar-10-plc':
        from data.cifar_plc import CIFAR10
        dataset = CIFAR10(root=p['data_path'], split='train', train_ratio=0.9, trust_ratio=0, download=True, transform=transform)

    elif p['train_db_name'] == 'cifar-100-plc':
        from data.cifar_plc import CIFAR100
        dataset = CIFAR100(root=p['data_path'], split='train', train_ratio=0.9, trust_ratio=0, download=True, transform=transform)
    else:
        raise ValueError('Invalid train dataset {}'.format(p['train_db_name']))

    # Wrap into other dataset (__getitem__ changes)

    if to_noisy_dataset:
        from data.custom_dataset import NoisyDataset
        dataset = NoisyDataset(dataset, meta_info)

    if to_augmented_dataset: # Dataset returns an image and an augmentation of that image.
        from data.custom_dataset import AugmentedDataset
        dataset = AugmentedDataset(dataset)

    if to_neighbors_dataset: # Dataset returns an image and one of its nearest neighbors.
        from data.custom_dataset import NeighborsDataset
        indices = np.load(p['topk_neighbors_train_path'])
        dataset = NeighborsDataset(dataset, indices, p['num_neighbors'], meta_info.get('predicted_labels'))

    return dataset

def get_val_dataset(p, transform=None, to_neighbors_dataset=False, meta_info=None):
    # Base dataset

    if p['val_db_name'] == 'cifar-10':
        if p['setup'] in [ 'propmix']:
            from data.cifar_propmix import cifar_dataset
            dataset = cifar_dataset(dataset=p['dataset'],root_dir=p['data_path'],transform=transform, meta_info=meta_info)
        else:
            from data.cifar import CIFAR10
            dataset = CIFAR10(root=p['data_path'], train=False, transform=transform, download=True)

    elif p['val_db_name'] == 'cifar-100':
        if p['setup'] in ['propmix']:
            from data.cifar_propmix import cifar_dataset
            dataset = cifar_dataset(dataset=p['dataset'],root_dir=p['data_path'],transform=transform, meta_info=meta_info)

        else:
            from data.cifar import CIFAR100
            dataset = CIFAR100(root=p['data_path'], train=True, transform=transform, download=True)

    # elif p['val_db_name'] == 'stl-10':
    #     from data.stl import STL10
    #     dataset = STL10(split='test', transform=transform, download=True)

    # elif p['val_db_name'] == 'imagenet':
    #     from data.imagenet import ImageNet
    #     dataset = ImageNet(split='val', transform=transform)

    #elif p['val_db_name'] in ['imagenet_50', 'imagenet_100', 'imagenet_200']:
    elif str(p['val_db_name']) in ['imagenet_100']:
        from data.imagenet import ImageNetSubset

        # subset_file = './data/imagenet_subsets/%s.txt' %(p['val_db_name'])
        #dataset = ImageNetSubset(subset_file=subset_file, split='val', transform=transform)
        dataset = ImageNetSubset(root_dir=p['data_path'], split='val', transform=transform)

    elif p['val_db_name'] == 'webvision':
        from data.webvision import Webvision
        dataset = Webvision(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'])
    elif p['val_db_name'] == 'mini_imagenet_blue':
        from data.redblue import MiniImagenet
        dataset = MiniImagenet(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'], color='blue')
    elif p['val_db_name'] in ['mini_imagenet_red','mini_imagenet32_red']:
        from data.redblue import MiniImagenet
        dataset = MiniImagenet(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'], color='red')
    elif p['val_db_name'] == 'stanford_cars_blue':
        from data.redblue import StanfordCars
        dataset = StanfordCars(root_dir=p['data_path'], transform=transform, meta_info=meta_info, num_classes=p['num_classes'], color='blue')
    elif p['val_db_name'] == 'cifar-10-plc':
        from data.cifar_plc import CIFAR10
        dataset = CIFAR10(root=p['data_path'], split='test', train_ratio=0.9, trust_ratio=0, download=True, transform=transform)

    elif p['val_db_name'] == 'cifar-100-plc':
        from data.cifar_plc import CIFAR100
        dataset = CIFAR100(root=p['data_path'], split='test', train_ratio=0.9, trust_ratio=0, download=True, transform=transform)
    else:
        raise ValueError('Invalid validation dataset {}'.format(p['val_db_name']))

    # Wrap into other dataset (__getitem__ changes)
    if to_neighbors_dataset: # Dataset returns an image and one of its nearest neighbors.
        from data.custom_dataset import NeighborsDataset
        indices = np.load(p['topk_neighbors_val_path'])
        dataset = NeighborsDataset(dataset, indices, 5) # Only use 5

    return dataset

def get_train_dataloader(p, dataset, shuffle=True, explicit_batch_size=None):
    batch_size = p['batch_size']
    if explicit_batch_size is not None:
        batch_size = explicit_batch_size
    return torch.utils.data.DataLoader(dataset, num_workers=p['num_workers'],
            batch_size=batch_size, pin_memory=True, collate_fn=collate_custom,
            drop_last=True, shuffle=shuffle)

def get_val_dataloader(p, dataset):
    return torch.utils.data.DataLoader(dataset, num_workers=p['num_workers'],
            batch_size=p['batch_size'], pin_memory=True, collate_fn=collate_custom,
            drop_last=False, shuffle=False)

def get_train_transformations(p):
    if p['augmentation_strategy'] == 'standard':
        # Standard augmentation strategy
        return transforms.Compose([
            transforms.RandomResizedCrop(**p['augmentation_kwargs']['random_resized_crop']),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(**p['augmentation_kwargs']['normalize'])
        ])

    elif p['augmentation_strategy'] == 'simclr':
        # Augmentation strategy from the SimCLR paper
        return transforms.Compose([
            transforms.RandomResizedCrop(**p['augmentation_kwargs']['random_resized_crop']),
            transforms.RandomHorizontalFlip(),
            transforms.RandomApply([
                transforms.ColorJitter(**p['augmentation_kwargs']['color_jitter'])
            ], p=p['augmentation_kwargs']['color_jitter_random_apply']['p']),
            transforms.RandomGrayscale(**p['augmentation_kwargs']['random_grayscale']),
            transforms.ToTensor(),
            transforms.Normalize(**p['augmentation_kwargs']['normalize'])
        ])

    elif p['augmentation_strategy'] == 'ours':
        # Augmentation strategy from our paper
        return transforms.Compose([
            transforms.RandomHorizontalFlip(),
            transforms.RandomCrop(p['augmentation_kwargs']['crop_size']),
            Augment(p['augmentation_kwargs']['num_strong_augs']),
            transforms.ToTensor(),
            transforms.Normalize(**p['augmentation_kwargs']['normalize']),
            Cutout(
                n_holes = p['augmentation_kwargs']['cutout_kwargs']['n_holes'],
                length = p['augmentation_kwargs']['cutout_kwargs']['length'],
                random = p['augmentation_kwargs']['cutout_kwargs']['random'])])

    #elif p['augmentation_strategy'] in  ['dividemix','dividemix_webv']:
    elif p['augmentation_strategy'] in  ['dividemix','dividemix_webv']:
        trnfs = []
        if 'resize' in p['augmentation_kwargs'].keys():
            trnfs = [transforms.Resize(p['augmentation_kwargs']['resize'])]
        try:
            padding = p['augmentation_kwargs']['crop_padding']
        except:
            padding = None
        trnfs.append(transforms.RandomCrop(p['augmentation_kwargs']['crop_size'], padding=padding))
        trnfs.append(transforms.RandomHorizontalFlip())
        trnfs.append(transforms.ToTensor())
        trnfs.append(transforms.Normalize(**p['augmentation_kwargs']['normalize']))
        return transforms.Compose(trnfs)
    elif p['augmentation_strategy'] == 'dividemix_webvision':
        trnfs = []
        if 'resize' in p['augmentation_kwargs'].keys():
            trnfs = [transforms.Resize(p['augmentation_kwargs']['resize'])]
        # try:
        #     padding = p['augmentation_kwargs']['crop_padding']
        # except:
        #     padding = None
        trnfs = [transforms.Resize(p['augmentation_kwargs']['resize'])]
        trnfs.append(transforms.RandomResizedCrop(p['augmentation_kwargs']['crop_size']))
        # trnfs.append(transforms.RandomCrop(p['augmentation_kwargs']['crop_size'], padding=padding))
        trnfs.append(transforms.RandomHorizontalFlip())
        trnfs.append(transforms.ToTensor())
        trnfs.append(transforms.Normalize(**p['augmentation_kwargs']['normalize']))
        return transforms.Compose(trnfs)
    elif p['augmentation_strategy'] == 'dividemix_red_mini_imagenet':
        trnfs = []

        # trnfs = [transforms.Resize(p['augmentation_kwargs']['random_resized_crop'])]

        trnfs.append(transforms.RandomResizedCrop(**p['augmentation_kwargs']['random_resized_crop']))
        # trnfs.append(transforms.RandomCrop(p['augmentation_kwargs']['crop_size'], padding=padding))
        trnfs.append(transforms.RandomHorizontalFlip())
        trnfs.append(transforms.ToTensor())
        trnfs.append(transforms.Normalize(**p['augmentation_kwargs']['normalize']))
        return transforms.Compose(trnfs)

    else:
        raise ValueError('Invalid augmentation strategy {}'.format(p['augmentation_strategy']))

def get_val_transformations(p):
    trnfs = []
    if 'resize' in p['augmentation_kwargs'].keys():
        trnfs = [transforms.Resize(p['augmentation_kwargs']['resize'])]
    # if p['augmentation_strategy'] != 'dividemix_red_mini_imagenet':
    #     trnfs.append(transforms.CenterCrop(p['augmentation_kwargs']['crop_size']))
    #elif p['augmentation_strategy'] == 'dividemix_webv':
    if p['augmentation_strategy'] == 'dividemix_webv':
        trnfs = [transforms.Resize(256)]
        trnfs.append(transforms.CenterCrop(p['augmentation_kwargs']['crop_size']))
    trnfs.append(transforms.ToTensor())
    trnfs.append(transforms.Normalize(**p['augmentation_kwargs']['normalize']))
    return transforms.Compose(trnfs)

def get_scan_transformations(p):
    if 'scan_kwargs' not in p.keys():
        return get_train_transformations(p)
    trnfs = [
        transforms.RandomHorizontalFlip(),
        transforms.RandomCrop(p['scan_kwargs']['crop_size']),
        Augment(p['scan_kwargs']['num_strong_augs']),
        transforms.ToTensor(),
        transforms.Normalize(**p['scan_kwargs']['normalize']),
        Cutout(
            n_holes = p['scan_kwargs']['cutout_kwargs']['n_holes'],
            length = p['scan_kwargs']['cutout_kwargs']['length'],
            random = p['scan_kwargs']['cutout_kwargs']['random'])]
    if 'resize' in p['scan_kwargs'].keys():
        trnfs.insert(0, transforms.Resize(p['scan_kwargs']['resize']))
    return transforms.Compose(trnfs)

def get_optimizer(p, model, cluster_head_only=False):
    if cluster_head_only: # Only weights in the cluster head will be updated
        for name, param in model.named_parameters():
                if 'cluster_head' in name:
                    param.requires_grad = True
                else:
                    param.requires_grad = False
        params = list(filter(lambda p: p.requires_grad, model.parameters()))
        assert(len(params) == 2 * p['num_heads'])

    else:
        params = model.parameters()

    if p['optimizer'] == 'sgd':
        optimizer = torch.optim.SGD(params, **p['optimizer_kwargs'])

    elif p['optimizer'] == 'adam':
        optimizer = torch.optim.Adam(params, **p['optimizer_kwargs'])

    else:
        raise ValueError('Invalid optimizer {}'.format(p['optimizer']))

    return optimizer

def adjust_learning_rate(p, optimizer, epoch):
    lr = p['optimizer_kwargs']['lr']

    if p['scheduler'] == 'cosine':
        eta_min = lr * (p['scheduler_kwargs']['lr_decay_rate'] ** 3)
        lr = eta_min + (lr - eta_min) * (1 + math.cos(math.pi * epoch / p['epochs'])) / 2

    elif p['scheduler'] == 'step':
        steps = np.sum(epoch > np.array(p['scheduler_kwargs']['lr_decay_epochs']))
        if steps > 0:
            lr = lr * (p['scheduler_kwargs']['lr_decay_rate'] ** steps)

    elif p['scheduler'] == 'constant':
        lr = lr

    else:
        raise ValueError('Invalid learning rate schedule {}'.format(p['scheduler']))

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

    return lr