#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jan 24 08:47:54 2018

@author: mohit

"""

import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
from tqdm import tqdm #to show progress
from sklearn.model_selection import train_test_split
from keras.models import Model
from keras.layers import Conv2D, MaxPool2D, Input
from sklearn.metrics import confusion_matrix, classification_report, f1_score
import xgboost as xgb

#Define Image Directory
image_dir_test = 'images/plants/test/'
image_dir_train = 'images/plants/train/'

#define the senstivity to control the selection of green color
sensitivity = 30
#define final image size for processing
image_size = 64
'''
define a function to remove background from the image to only leave the green leaves, blurring and normalizing. Followed by resizing the images
to 64 x 64 size

'''

def image_transformation(imageName, sensitivity):

    imagePlatCV = cv2.imread(imageName) #read image
    hsvImage = cv2.cvtColor(imagePlatCV, cv2.COLOR_BGR2HSV)
    #define the range for green color
    lower_green = np.array([60 - sensitivity, 100, 50])
    upper_green = np.array([60 + sensitivity, 255, 255])
    # threshold the hsv image to get only green colors
    mask = cv2.inRange(hsvImage, lower_green, upper_green)
    #apply bitwise_and between mask and the original image
    greenOnlyImage = cv2.bitwise_and(imagePlatCV, imagePlatCV, mask=mask)
    #lets define a kernal with ones
    kernel0 = np.ones((15,15), np.uint8)
    #lets apply closing morphological operation
    closing0 = cv2.morphologyEx(greenOnlyImage, cv2.MORPH_CLOSE, kernel0)
    #blur the edges
    blurImage = cv2.GaussianBlur(closing0, (15,15), 0)
    blurImageColor = cv2.cvtColor(blurImage, cv2.COLOR_BGR2RGB)#to make it work with right color
    #resize image
    resizeImage = cv2.resize(blurImageColor, (image_size, image_size), interpolation=cv2.INTER_AREA)
    resizeImage = resizeImage/255 #normalize
    #resizeImage = resizeImage.reshape(image_size,image_size,3) #to make it in right dimensions for the Keras add 3 channel
    print(resizeImage.shape)
    return resizeImage

#define list of plant species
classes = ['Black-grass', 'Charlock', 'Cleavers', 'Common Chickweed', 'Common wheat', 'Fat Hen', 'Loose Silky-bent', 'Maize', 'Scentless Mayweed'
           , 'Shepherds Purse', 'Small-flowered Cranesbill', 'Sugar beet']
'''
Data extraction: The loop below will create a data list containing image file path name, the classifcation lable
(0 -11) and the specific plant name

'''
train = [] #data list
for species_lable, speciesName in enumerate(classes):
    for fileName in os.listdir(os.path.join(image_dir_train, speciesName)):
        train.append([image_dir_train + '{}/{}'.format(speciesName, fileName), species_lable, speciesName])

#convert the list into dataframe using Pandas
trainigDataFrame = pd.DataFrame(train, columns=['FilePath', 'PlantLabel', 'PlantName'])

#Suffle the data
seed = 1234 #define seed to get consistent results
trainigDataFrame = trainigDataFrame.sample(frac=1, random_state=seed)
trainigDataFrame = trainigDataFrame.reset_index()

#Prepare the images for the model by preprocessing

X = np.zeros((trainigDataFrame.shape[0], image_size, image_size, 3)) #array to store image after image_transformfunction

for i, fileName in tqdm(enumerate(trainigDataFrame['FilePath'])):
    print(fileName)
    newImage = image_transformation(fileName, sensitivity)
    X[i] = newImage

#Convert lables to categorical and do one-hot encoding
y = trainigDataFrame['PlantLabel']
y = np.array(y)
y_labels_cat, _ = pd.factorize(y)

'''
Build the initial network from pretrained model that will be used for transfer learning. We will use first two blocks from
VGG19 pretrained model
'''
#input layer
imageInput = Input(shape=(image_size, image_size, 3), name='input_1')
#Block 1 - layers name same as the layers in pretrained model
model = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv1')(imageInput)
model = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv2')(model)
model = MaxPool2D(pool_size= (2,2), strides=(2,2), name='block1_pool')(model)

#Block 2

model = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv1')(model)
model = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv2')(model)
model = MaxPool2D(pool_size= (2,2), strides=(2,2), name='block2_pool')(model)
model_initial = Model(inputs=imageInput, outputs=model)
model_initial.summary()

#define the dictionary of layers
layer_dict = dict([(layer.name, layer) for layer in model_initial.layers])
print(layer_dict)

#load weights from VGG19 model
weights_path = 'vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5'#from https://github.com/fchollet/deep-learning-models/releases/
model_initial.load_weights(weights_path, by_name=True)

#set weights for initial layers
for i, layer in enumerate(model_initial.layers):
    weights = layer.get_weights()
    model_initial.layers[i].set_weights(weights)

model_final = Model(inputs = model_initial.input, outputs = model_initial.output)

features = [] #List to extract features using pretrained model weights
for i in tqdm(X):
    image = np.expand_dims(i, axis=0)
    featurePredict = model_final.predict(image)
    features.append(featurePredict)

#convert to numpy array from list
featuresArray = np.array(features)

#reshape to get in the right shape for keras model
featuresReshape = np.reshape(featuresArray, (featuresArray.shape[0], featuresArray.shape[2], featuresArray.shape[3], featuresArray.shape[4]))

featuresReshapeXGboost = np.reshape(featuresReshape, (featuresReshape.shape[0], featuresReshape.shape[1]*featuresReshape.shape[2]*featuresReshape.shape[3]))

print(featuresReshapeXGboost[1].shape)

#split dataset into Train and Test
X_train, X_val, y_train, y_val = train_test_split(featuresReshapeXGboost, y_labels_cat, test_size=0.20, random_state=seed)

#Use xgboost from SKLearn Library as classifier on top of the model build using transfer learning

xgb1 = xgb.sklearn.XGBClassifier(
        learning_rate =0.1,
         n_estimators=100,
         max_depth=5,
         min_child_weight=11,
         gamma=0.1,
         subsample=0.8,
         colsample_bytree=0.7,
         objective='multi:softprob',
         n_jobs=-1,
         scale_pos_weight=1,
         seed=seed)

xgb1.fit(X_train, y_train)

#Predict the classes

y_pred = xgb1.predict(X_val)

#confusion matrix and classification report

cm = confusion_matrix(y_val, y_pred)

print(classification_report(y_val, y_pred, target_names=classes))

#plot F1-score vs. Classes

f1Score = f1_score(y_val, y_pred, average=None)

y_pos = np.arange(len(classes))
plt.bar(y_pos, f1Score)
plt.xticks(y_pos, classes)
plt.ylabel('F1 Score')
plt.title('F1 Score of various species after classification')
plt.show()


'''
Twitter_Presidential_Race_Sentiment_Clustering: ProcessTwitterData.py
Authors: Justin Murphey, Rafael Zamora
Last Updated: 12/04/16

CHANGE-LOG:
-Code Review/Refactoring
-Updated Comments

'''

"""
This script is used to process data gathered from Twitter
and generate a data set which will be used in the cluster analysis.

"""

import os, json, re
import numpy as np
from textblob import TextBlob
from textblob.sentiments import NaiveBayesAnalyzer
from difflib import SequenceMatcher
from multiprocessing import Pool

path_to_data = "/Twitter_Presidential_Race_Sentiment_Clustering/data/"

def process_tweet_data(tweets_data):
    '''
    process_tweet_data() calculates the candidate reference value and sentiment
    value of tweet data. Candidate reference value is calulated by taking the
    difference of the normalized noun-phrases/keyword similarity value of both candidates.
    Sentiment value is calculated using NaiveBayes Classifier trained on movie review data.
    It then returns list of processed data.

    '''
    raw_ref_hil_datalist = []
    raw_ref_tru_datalist = []
    polarity_datalist = []
    analyzer_NB = NaiveBayesAnalyzer()
    for tweet_data in tweets_data:
        tweet_data = tweet_data.split(',')
        tweet_text = tweet_data[2]
        tweet_text = re.sub(r'pic.*$', '', tweet_text)
        tweet_text = re.sub(r'http.*$', '', tweet_text)
        tweet_text = tweet_text.replace("@", "")
        tweet_text = tweet_text.replace("#", "")
        tweet_text = tweet_text.replace("'", "")

        tweet_text_blob = TextBlob(tweet_text, analyzer=analyzer_NB)
        tweet_polarity = (tweet_text_blob.sentiment.p_pos*2.0) - 1.0
        tweet_ref_hil = max_noun_phrase_similarity(tweet_text, ["hillary clinton", "hillaryclinton"])
        tweet_ref_tru = max_noun_phrase_similarity(tweet_text, ["donald trump", "realdonaldtrump"])
        raw_ref_hil_datalist.append(tweet_ref_hil)
        raw_ref_tru_datalist.append(tweet_ref_tru)
        polarity_datalist.append(tweet_polarity)

    max_value_hil = max(raw_ref_hil_datalist)
    max_value_tru = max(raw_ref_tru_datalist)
    min_value_hil = min(raw_ref_hil_datalist)
    min_value_tru = min(raw_ref_tru_datalist)
    norm_ref_datalist = []
    for i in range(len(raw_ref_hil_datalist)):
        norm_ref = float(raw_ref_hil_datalist[i]-min_value_hil/max_value_hil - min_value_hil)-float(raw_ref_tru_datalist[i]-min_value_tru/max_value_tru - min_value_tru)
        norm_ref_datalist.append(norm_ref)
    processed_tweets_data = []
    for i in range(len(polarity_datalist)): processed_tweets_data.append([polarity_datalist[i], norm_ref_datalist[i]])
    return np.array(processed_tweets_data)

def max_noun_phrase_similarity(text, keywords):
    '''
    max_noun_phrase_similarity() finds string similarity between the keywords
    and the noun phrases of the text. It then returns the highest similarity value.

    '''
    noun_phrases = TextBlob(text).noun_phrases
    print(noun_phrases)
    values = []
    for keyword in keywords:
        for nf in noun_phrases: values.append(SequenceMatcher(None, keyword, nf).ratio())
    if len(values) == 0: return 0
    else: return max(values)

if __name__ == '__main__':
    ''''
    Parameters:
    parallel -set True to run distributed processing
    cores -number of cores used for distributed processing

    ''''
    parallel = False
    cores = 2

    '''
    Process all raw data files in data directory.

    '''
    for filename in sorted(os.listdir(path_to_data+"raw/")):
        tweets_data = []
        print("Processing: "+ filename)
        with open(path_to_data + filename, "r") as f: tweets_data = f.readlines()

        if parallel:
            tweets_data_parts = np.array_split(np.array(tweets_data),cores)
            tweets_data_parts = [sublist.tolist() for sublist in tweets_data_parts]
            with Pool() as pool:
                results = pool.map(process_tweet_data, tweets_data_parts)
            dataset = [val for sublist in results for val in sublist]
        else:
            dataset = process_tweet_data(tweets_data)

        os.makedirs(os.path.dirname(path_to_data+"processed/processed_"+filename), exist_ok=True)
        np.savetxt(path_to_data+"processed/processed_"+filename, dataset, fmt='%.5f', delimiter=",")
        print("Processing Done.\nProcessed File: processed_"+filename)


#! /usr/bin/env python

from sys import argv
import numpy as np
from ase.io import read, write
from ase.build.surface import fcc111

if argv[1] == "-h" or argv[1] == "--help":
    print("""
Python script that generates a molecular junction for transport computation
using DFTB+.

Author: Alessandro Pirrotta
MIT Licence

Example 1:
$ python generate_molecular_junction.py molecule.xyz c r 0 1

using molecule 'molecule.xyz' as device region
it generates a cluster type molecular junction (c) with no pbc,
removes terminal hydrogen (r) on the end atoms 0 and 1.

Example 2:
$ python generate_molecular_junction.py molecule.xyz s nr

generate a supercell type molecular junction with pbc (s),
does not removes terminal hydrogen (nr) and automatically looks for
sulphur atoms as end atoms.
""")
    exit()

def find_nearest(atom, atoms):
    # create empty list
    dists = []
    # define coordinates for atom
    coord = atom.position
    # loop over atoms and append distance to list
    for a in atoms:
        if a.symbol == 'H':
            v = a.position - coord
            dist = np.sqrt(v[0]**2+v[1]**2+v[2]**2)
            dists.append(dist)
            continue
        dists.append(1.7)
    # make list to np.array
    dists = np.array(dists)
    # find index with shortest distance using argmin
    index = np.argmin(dists)
    if dists[index] > 1.7:  # atom closer than C
        print 'found no H bonded to S, deleted something, consider revising'
    return index

if len(argv) > 4:
    # print("reading end atoms from input")
    endatom1 = int(argv[4])
    endatom2 = int(argv[5])

# create gold electrode
a = 4
pbc = bool(True)

b = a
c = 6
lelectrode = fcc111('Au', size=(a, b, c), vacuum=5.0, orthogonal=True)

# Au-S 2.72AA / Au-Au 2.88AA
# distance between two adjacent gold layers
# deltag           2.35558909829
# perpendicular distance between endatom and electrode surface
# delta s_g        2.14441090171

# read the molecule file
# molecule = read('mol.xyz')
molecule = read(argv[1])
# place electrode close to the origin of axis
n = (lelectrode[0].position)
n = - n                                 # put atom 1 at the origin of axis
lelectrode.translate(n)
natoms = len(molecule)

try:
    endatom1
except NameError:
    endatom1 = -1
    for atom in molecule:
        # if your endatoms are S atoms and there are only 2 in the molecule
        if atom.symbol == 'S':
            if endatom1 < 0:
                endatom1 = atom.index
            else:
                endatom2 = atom.index

# remove terminal hydrogens
if len(argv) > 3 and argv[3] == 'r':
    list = []
    for atom in molecule:
        if atom.symbol == 'S':
            list.append(atom.index)
    del molecule[find_nearest(molecule[list[1]], molecule)]
    del molecule[find_nearest(molecule[list[0]], molecule)]
    natoms = len(molecule)

# align the sulphur atoms along the z axis
po = (molecule[endatom1].position)
lo = (molecule[endatom2].position)
v = lo-po
z = [0, 0, 1]
molecule.rotate(v, z)
molecule.rotate('z', np.pi+np.pi/4.0)

# place the molecule on top of a hollow fcc site for axa face
if a == 6:
    q = (lo + [0, 0, 4.5] - lelectrode[(b*a*2)+(6*4)].position)
if a == 3:
    q = (lo + [0, 0, 4.5] - lelectrode[(a*b*2)+(4)].position)
if a == 4:
    q = (lo + [0, 0, 4.5] - lelectrode[(b*a*2)+(9)].position)
q = - q
molecule.translate(q)

# define distance betweet gold layers
deltag = abs(lelectrode[0].position[2] - lelectrode[(a*b)+1].position[2])
# define end to end distance
v = molecule[endatom2].position - molecule[endatom1].position
lenght_molecule = np.sqrt(v[0]**2+v[1]**2+v[2]**2)
molecule.translate([0, 0, -(deltag)])
molecule.extend(lelectrode)
# make a list of indexes relative to gold atoms
au_list = []
for atom in molecule:
    if atom.symbol == 'Au':
        au_list.append(atom.index)

# number of gold atoms on a single layer
ng = int(a*b)
# distance along the z axis between one layer and the adjacent one:
deltag = abs(lelectrode[0].position[2] - lelectrode[(a*b)+1].position[2])
v = molecule[endatom2].position-molecule[natoms+2].position
deltas_g = abs(v[2])

for x in range(c):
    ly = []
    ly = au_list[(x*a*b):(a*b*(x+1))]
    layer = molecule[ly]
    layer.translate([0, 0,
                    -(abs(lenght_molecule)+abs(deltas_g*2)+abs(deltag*(2*x)))])
    molecule.extend(layer)

if argv[2] == "s" or argv[2] == "S":
    molecule.set_pbc(True)
if argv[2] == "c" or argv[2] == "C":
    molecule.set_pbc(False)
n = (molecule[natoms+a*b*c-a*b].position)
n = -n
molecule.translate(n)

if argv[2] == "s" or argv[2] == "S":
    lele = fcc111('Au',
                  size=(a, b, 18),
                  vacuum=deltas_g+lenght_molecule,
                  orthogonal=True)
    cell = lele.get_cell()
    molecule.set_cell([cell[0],
                       cell[1],
                       [0., 0., molecule[len(molecule)-1].position[2]-20.0]])
    molecule.set_pbc((True, True, True))
    write('moljunct_S.gen', molecule)

write('moljunct_C.gen', molecule)

# printing device, source and drain atom indeces
print 1, natoms
print natoms+1, natoms+a*b*6
print natoms+a*b*6+1, len(molecule)

exit()


import time
from os import path
from random import SystemRandom, seed, randint

import cv2
import numpy as np
import wandb
from torch import optim, save
from torch.multiprocessing import Process

from serpentrain.reinforcement_learning.distributed.utils import create_model, create_controller, create_env

class ModelTrainServer(Process):
    """
    Pulls samples from sample queue
    Train Model
    Share model via shared model dict
    """

    def __init__(self, sample_queue, model_shared_dict, state_dict_optimizer, save_dir, device="cuda:0",
                 episode_length=1000, num_eval_episodes=10, save_interval=30):
        super().__init__()
        self.save_interval = save_interval
        self.sample_queue = sample_queue
        self.device = device
        self.model_shared_dict = model_shared_dict
        self.save_dir = save_dir
        self.model = create_model(state_dict=model_shared_dict, device=device)
        self.episode_length = episode_length
        self.num_eval_episodes = num_eval_episodes
        self.state_dict_optimizer = state_dict_optimizer

        self.gamma = 0.99
        self.steps = 10_000

        self.shutdown = False

    def _initialize(self):
        # Initialize wandb
        wandb.init(project="flatland")

        # Create the controller so we can call the specific train function
        # The controller will not be used for anything else, therefore the model can be anything
        self.controller = create_controller(self.model)

        self.optimizer = optim.Adam(self.model.parameters(), lr=0.5e-4)

        if self.state_dict_optimizer:
            self.optimizer.load_state_dict(self.state_dict_optimizer)

    def run(self) -> None:
        self._initialize()
        time_last_save = time.time()  # in seconds
        epoch = 0
        wandb.watch(self.model)
        while not self.shutdown:
            print(f"Starting Training on {self.device}")
            start_time = time.time()
            self.train(epoch)
            print(f"Training Episode finished in {time.time() - start_time}")
            self.update_shared_model()

            print(f"Evaluate")
            self.evaluate()
            print(f"Finished Evaluating")

            time_current = time.time()
            if time_current - time_last_save >= self.save_interval * 60:
                time_last_save = time_current
                name = f"snapshot-{time.strftime('%Y%m%d-%H%M')}-epoch-{epoch}.pt"
                save({
                    "model": self.model.state_dict(),
                    "optimizer": self.optimizer.state_dict()
                }, path.join(self.save_dir, name))
                print("Saved Snapshot")
            epoch += 1

    def train(self, epoch):
        """
        Does self.steps training steps
        """
        losses = []
        for i in range(self.steps):
            losses.append(self.train_step(i))
        total_loss = sum(losses)
        wandb.log({"Epoch Loss": total_loss / len(losses), "Epoch": epoch}, commit=False)

    def train_step(self, i):
        """
        Does a single batch training step
        """
        loss = self.controller.train(self.model, self, self.gamma, optimizer=self.optimizer)
        wandb.log({"Training loss": loss, "Step": i}, commit=False)
        return loss

    def evaluate(self):
        # Set a random state
        rng = SystemRandom()
        seed_ = rng.randint(0, 2 ** 32 - 1)
        seed(seed_)
        np.random.seed(seed_)

        # Create and Start Environment
        env = create_env(seed=seed_)

        # renderer = RenderTool(env, gl="PIL")

        # model_ref = ModelWrapper.wrap_model(self.backlog_queue, self.return_dict, self.id)
        self.controller.model = self.model

        env.number_of_agents = randint(1, 100)
        env.width = randint(10, 100)
        env.height = randint(10, 100)

        sum_score = 0
        sum_score_per_agent = 0
        sum_steps = 0
        sum_agents = 0
        sum_time = 0
        sum_agents_done = 0
        sum_agents_done_percentage = 0
        total_action_dict = {0: 0,
                             1: 0,
                             2: 0,
                             3: 0,
                             4: 0}

        for game in range(self.num_eval_episodes):
            start_time = time.time()
            print(f"Evaluating for {game} out of {self.num_eval_episodes} games.")
            obs, info = env.reset(regenerate_rail=True, regenerate_schedule=True, random_seed=True)
            # renderer.reset()
            # images = [renderer.get_image()]
            score = 0
            # trajectories = [Trajectory() for _ in env.get_agent_handles()] # Not needed

            # Create and Start Controller
            self.controller.start_of_round(obs=obs, env=env)

            step = 0
            done = {}

            for step in range(self.episode_length):
                action_dict, processed_obs = self.controller.act(observation=obs, info=info, epsilon=0)
                obs, all_rewards, done, info = env.step(action_dict)

                for action in action_dict.values():
                    total_action_dict[action] = total_action_dict.get(action, 0) + 1

                score += sum(all_rewards)

                if done['__all__'] or self.shutdown:
                    break

            n_agents_done = sum(done.values()) - 1
            n_agents = env.get_num_agents()
            score_per_agent = score / n_agents
            agents_done_percentage = n_agents_done / n_agents

            sum_score += score
            sum_score_per_agent += score_per_agent
            sum_steps += step
            sum_agents += n_agents
            sum_agents_done += n_agents_done
            sum_agents_done_percentage += agents_done_percentage
            sum_time = time.time() - start_time

            print(f"Finished evaluation game #{game}.\n"
                  f"Score: {score}. #Agents: {n_agents}.\n"
                  f"Score per agent: {score_per_agent}."
                  f"Steps: {step}. Time: {time.time() - start_time}"
                  f"#Agents done: {n_agents_done}. %Agents done: {agents_done_percentage}.")

        avg_score = sum_score / self.num_eval_episodes
        avg_score_per_agent = sum_score_per_agent / self.num_eval_episodes
        avg_steps = sum_steps / self.num_eval_episodes
        avg_agents = sum_agents / self.num_eval_episodes
        avg_agents_done = sum_agents_done / self.num_eval_episodes
        avg_agents_done_percentage = sum_agents_done_percentage / self.num_eval_episodes

        print("=" * 20)
        print(f"Finished all evaluation games. Averages:")
        print(f"Score: {avg_score}. #Agents: {avg_agents}. "
              f"Score per agent: {avg_score_per_agent}. Steps: {avg_steps}. "
              f"#Agents done: {avg_agents_done}. %Agents done: {avg_agents_done_percentage}.\n"
              f"Total Time: {sum_time} per episode")
        print("=" * 20)

        data = [[label, val] for (label, val) in total_action_dict.items()]
        table = wandb.Table(data=data, columns=["action", "count"])

        wandb.log({
            "Action Distribution": wandb.plot.bar(table, "action", "count", title="Action Distribution"),
            "Average Score": avg_score,
            "Average Agents": avg_agents,
            "Score per Agent": avg_score_per_agent,
            "Average Steps": avg_steps,
            "Average Agents Done": avg_agents_done,
            "Percentage Agents Done": avg_agents_done_percentage}, commit=True)

    def update_model(self):
        print("Updated model")
        if self.model is None:
            self.model = create_model(self.model_shared_dict, device=self.device)
        else:
            self.model.load_state_dict(self.model_shared_dict)
        self.model.eval()

    def update_shared_model(self):
        print("Updating global model")
        cpu_state_dict = {k: v.cpu() for k, v in self.model.state_dict().items()}
        self.model_shared_dict.update(cpu_state_dict)

    def sample(self):
        """
        Wrapper function for the controller train function
        """
        return self.sample_queue.get()

    @staticmethod
    def log_video(_images, episode):
        height, width, depth = _images[0].shape
        out = cv2.VideoWriter(f'flatland_{episode}.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
        [out.write(cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)) for image in _images]
        out.release()
        wandb.log({"Replay": wandb.Video(f"flatland_{episode}.mp4")}, commit=False)


#!/usr/bin/env python3
import argparse
import os
import time
from datetime import timedelta
import sys
import tensorflow as tf
import numpy as np
from seqdesign.version import VERSION
from seqdesign import hyper_conv_auto as model
from seqdesign import helper
from seqdesign import utils
from seqdesign import aws_utils

def main(working_dir='.'):
    start_run_time = time.time()

    parser = argparse.ArgumentParser(description="Train an autoregressive model on a collection of sequences.")
    parser.add_argument("--s3-path", type=str, default='',
                        help="Base s3:// path (leave blank to disable syncing).")
    parser.add_argument("--s3-project", type=str, default=VERSION, metavar='P',
                        help="Project name (subfolder of s3-path).")
    parser.add_argument("--run-name-prefix", type=str, default=None, metavar='P',
                        help="Prefix for run name.")
    parser.add_argument("--channels", type=int, default=48, metavar='C',
                        help="Number of channels.")
    parser.add_argument("--num-iterations", type=int, default=250005, metavar='N',
                        help="Number of iterations to run the model.")
    parser.add_argument("--snapshot-interval", type=int, default=None, metavar='N',
                        help="Take a snapshot every N iterations.")
    parser.add_argument("--dataset", type=str, default=None,
                        help="Dataset name for fitting model. Alignment weights must be computed beforehand.")
    parser.add_argument("--restore", type=str, default='',
                        help="Session name for restoring a model to continue training.")
    parser.add_argument("--gpu", type=str, default='',
                        help="Which gpu to use. Usually  0, 1, 2, etc...")
    parser.add_argument("--r-seed", type=int, default=42, metavar='RSEED',
                        help="Random seed for parameter initialization")
    parser.add_argument("--alphabet-type", type=str, default='protein', metavar='T',
                        help="Type of data to model. Options = [protein, DNA, RNA]")
    ARGS = parser.parse_args()

    if ARGS.gpu != '':
        os.environ["CUDA_VISIBLE_DEVICES"] = ARGS.gpu
    dataset_name = ARGS.dataset.rsplit('/', 1)[-1]
    if dataset_name.endswith(".fa"):
        dataset_name = dataset_name[:-len(".fa")]
    elif dataset_name.endswith(".fasta"):
        dataset_name = dataset_name[:-len(".fasta")]

    if ARGS.restore:
        # prevent from repeating batches/seed when restoring at intermediate point
        # script is repeatable as long as restored at same step
        # assumes restore arg of *.ckpt-(int step)
        restore_ckpt = ARGS.restore.split('.ckpt-')[-1]
        r_seed = ARGS.r_seed + int(restore_ckpt)
        r_seed = r_seed % (2 ** 32 - 1)  # limit of np.random.seed
    else:
        r_seed = ARGS.r_seed

    print(ARGS.restore)

    if ARGS.restore == '':
        folder_time = (
            f"{dataset_name}_{ARGS.s3_project}_channels-{ARGS.channels}"
            f"_rseed-{ARGS.r_seed}_{time.strftime('%y%b%d_%I%M%p', time.gmtime())}"
        )
        if ARGS.run_name_prefix is not None:
            folder_time = ARGS.run_name_prefix + '_' + folder_time
    else:
        folder_time = ARGS.restore.split('/')[-1]
        folder_time = folder_time.split('.ckpt')[0]

    folder = f"{working_dir}/sess/{folder_time}"
    os.makedirs(folder, exist_ok=True)
    log_f = utils.Tee(f'{folder}/log.txt', 'a')  # log stdout to log.txt

    print(folder)
    print(ARGS)

    print("OS: ", sys.platform)
    print("Python: ", sys.version)
    print("TensorFlow: ", tf.__version__)
    print("Numpy: ", np.__version__)

    print('Using device:', ARGS.gpu)
    available_devices = utils.get_available_gpus()
    if available_devices:
        print('\t'.join(available_devices))
        print('\t'.join(utils.get_available_gpus_desc()))
        print(utils.get_cuda_version())
        print("CuDNN Version ", utils.get_cudnn_version())

    print("SeqDesign git hash:", str(utils.get_github_head_hash()))
    print()

    print("Run:", folder_time)

    tf.set_random_seed(r_seed)

    aws_util = aws_utils.AWSUtility(s3_base_path=ARGS.s3_path, s3_project=ARGS.s3_project) if ARGS.s3_path else None
    data_helper = helper.DataHelperSingleFamily(
        working_dir=working_dir, dataset=ARGS.dataset,
        r_seed=r_seed, alphabet_type=ARGS.alphabet_type,
        aws_util=aws_util,
    )
    data_helper.family_name = dataset_name

    # Variables for runtime modification
    batch_size = 30
    if ARGS.snapshot_interval is not None:
        fitness_check = ARGS.snapshot_interval
        fitness_start = 1
    else:
        fitness_check = 25000
        fitness_start = 29999
    num_iterations = int(ARGS.num_iterations)
    N_pred_iterations = 50

    plot_train = 100
    dropout_p_train = 0.5

    # fitness_check = 2
    # fitness_start = 2
    # num_iterations = 36

    # dims = {}
    dims = {"alphabet": len(data_helper.alphabet)}
    conv_model = model.AutoregressiveFR(dims=dims, channels=ARGS.channels)

    params = tf.trainable_variables()
    p_counts = [np.prod(v.get_shape().as_list()) for v in params]
    p_total = sum(p_counts)
    print(f"Total parameter number: {p_total}\n")

    saver = tf.train.Saver()
    run_metadata = tf.RunMetadata()

    with tf.Session() as sess:
        # Initialization
        print('Initializing variables')
        init = tf.global_variables_initializer()
        sess.run(init)

        print('Parameters initialized')

        if ARGS.restore != '':
            sess_namedir = f"{folder}/{ARGS.restore}"
            saver.restore(sess, sess_namedir)
            print("Parameters restored")
        else:
            # save initialization state
            save_path = saver.save(sess, f"{folder}/{folder_time}.ckpt", global_step=0)

        # Summary output
        merged = tf.summary.merge_all()
        train_writer = tf.summary.FileWriter(folder, sess.graph)
        global_step = 0
        print('Created train writer')

        # Run optimization
        while global_step < num_iterations:
            start = time.time()

            prot_decoder_input_f, prot_decoder_output_f, prot_mask_decoder_f, \
            family_one_hot_f, Neff_f = data_helper.generate_one_family_minibatch_data(batch_size)

            prot_decoder_input_r, prot_decoder_output_r, prot_mask_decoder_r, \
            family_one_hot_r, Neff_r = data_helper.generate_one_family_minibatch_data(batch_size, reverse=True)

            data_load_time = time.time() - start

            feed_dict = {
                conv_model.placeholders["sequences_start_f"]: prot_decoder_input_f,
                conv_model.placeholders["sequences_stop_f"]: prot_decoder_output_f,
                conv_model.placeholders["mask_decoder_1D_f"]: prot_mask_decoder_f,
                conv_model.placeholders["Neff_f"]: [Neff_f],
                conv_model.placeholders["sequences_start_r"]: prot_decoder_input_r,
                conv_model.placeholders["sequences_stop_r"]: prot_decoder_output_r,
                conv_model.placeholders["mask_decoder_1D_r"]: prot_mask_decoder_r,
                conv_model.placeholders["Neff_r"]: [Neff_r],
                conv_model.placeholders["step"]: [global_step],
                conv_model.placeholders["dropout"]: dropout_p_train
            }

            summary, _, global_step, ce_loss, loss, KL_embedding_loss \
                = sess.run([merged, conv_model.opt_op, conv_model.global_step,
                            conv_model.tensors["cross_entropy_loss"], conv_model.tensors["loss"],
                            conv_model.tensors["KL_embedding_loss"]], feed_dict=feed_dict)

            print(f'{global_step:7} {time.time() - start:0.3f} {data_load_time:0.3f} '
                  f'{ce_loss:10.4f} {loss:10.4f} {KL_embedding_loss:10.4f}', flush=True)

            if global_step % plot_train == 0:
                train_writer.add_summary(summary, global_step)

            if global_step % fitness_check == 0 and global_step > fitness_start:
                save_path = saver.save(sess, f"{folder}/{folder_time}.ckpt", global_step=global_step)
                if working_dir != '.':
                    data_helper.generate_reset_script(
                        f"{folder_time}.ckpt-{global_step}", ARGS.channels, ARGS.dataset, ARGS.r_seed)
        try:
            max_gpu_mem_used = sess.run(tf.contrib.memory_stats.MaxBytesInUse())
        except tf.errors.OpError:
            max_gpu_mem_used = None

        train_writer.close()

    print(f"Done! Total run time: {timedelta(seconds=time.time()-start_run_time)}")
    log_f.flush()
    if aws_util:
        aws_util.s3_sync(local_folder=folder, s3_folder=f'sess/{folder_time}/', destination='s3')

    if working_dir != '.':
        os.makedirs(f'{working_dir}/complete/', exist_ok=True)
        OUTPUT = open(f'{working_dir}/complete/{folder_time}.txt', 'w')
        OUTPUT.write(f"STEPS COMPLETED: {int(global_step)}")
        if max_gpu_mem_used is not None:
            OUTPUT.write(f"Max GPU memory used: {max_gpu_mem_used}")
        OUTPUT.close()

if __name__ == "__main__":
    main()