import math
from datetime import datetime
from collections import defaultdict
import numpy as np
import h5py
import pandas as pd

from exetera.core import exporter, persistence, utils, session, fields
from exetera.core.persistence import DataStore
from exetera.processing.nat_medicine_model import nature_medicine_model_1
from exetera.processing.method_paper_model import method_paper_model

"""
Goal
 * show whether mask wearing helps with exposure
   * must separate out preventative / remedial action
     * is the person sick and if so, when
     * has the person been predicted as having covid and if so, when
     * has the person had a positive covid test result and if so, when
   * preventative
     * healthy
     * not healthy but not yet covid
   * remedial
     * having had covid
   * must separate out confounding factors
     * background risk
       * non-health worker / non-covid health worker / covid health worker

Process
1. generate results for tested people
  1. filter out people without tests

n. generate results for predicted people
"""

def ppe_use_and_travel(ds, src, tmp, start_timestamp):

    logging = True

    s_asmts = src['assessments']

    if 'filtered_assessments' not in tmp.keys():
        f_asmts = tmp.create_group('filtered_assessments')
        cats = ds.get_reader(s_asmts['created_at'])
        asmt_filter = cats[:] >= start_timestamp

        ccs = ds.get_reader(s_asmts['country_code'])
        asmt_filter = asmt_filter & (ccs[:] == b'GB')

        symptom_keys = ('persistent_cough', 'fatigue', 'delirium', 'shortness_of_breath',
                        'fever', 'diarrhoea', 'abdominal_pain', 'chest_pain', 'hoarse_voice',
                        'skipped_meals', 'loss_of_smell')
        mask_keys = ('mask_cloth_or_scarf', 'mask_surgical', 'mask_n95_ffp')
        isolation_keys = ('isolation_healthcare_provider', 'isolation_little_interaction',
                          'isolation_lots_of_people')
        other_keys = ('patient_id', )
        symptom_thresholds = {s: 2 for s in symptom_keys}
        symptom_thresholds.update({m: 2 for m in mask_keys})
        symptom_thresholds['fatigue'] = 3
        symptom_thresholds['shortness_of_breath'] = 3

        for k in symptom_keys + mask_keys + isolation_keys + other_keys:
            with utils.Timer("filtering {}".format(k)):
                reader = ds.get_reader(s_asmts[k])
                if k in mask_keys + symptom_keys:
                    values = np.where(reader[:] >= symptom_thresholds[k], 1, 0)
                    ds.get_numeric_writer(f_asmts, k, 'int8').write(
                        ds.apply_filter(asmt_filter, values))
                    hist = np.unique(reader[:], return_counts=True)
                    print(sorted(zip(hist[0], hist[1])))
                    hist = np.unique(values, return_counts=True)
                    print(sorted(zip(hist[0], hist[1])))
                else:
                    reader.get_writer(f_asmts, k).write(ds.apply_filter(asmt_filter, reader))

        print('filtered assessments:', np.count_nonzero(asmt_filter), len(asmt_filter))
    #
    #
    # if 'filtered_assessment_predictions' not in tmp.keys():
    #     f_pred_asmts = tmp.create_group('filtered_assessment_predictions')
        symptom_readers = dict()
        for s in symptom_keys:
            symptom_readers[s] = ds.get_reader(f_asmts[s])
        predictions = ds.get_numeric_writer(f_asmts, 'prediction', 'float32')
        method_paper_model(ds, symptom_readers, predictions)
        predictions = ds.get_reader(f_asmts['prediction'])
        print('predictions:', np.count_nonzero(predictions[:] > 0), len(predictions))

    if 'patient_assessment_summaries' not in tmp.keys():
        asmt_psum = tmp.create_group('patient_assessment_summaries')
        pids = ds.get_reader(f_asmts['patient_id'])
        mcos = ds.get_reader(f_asmts['mask_cloth_or_scarf'])
        msurg = ds.get_reader(f_asmts['mask_surgical'])
        m95 = ds.get_reader(f_asmts['mask_n95_ffp'])
        with utils.Timer("generating patient_id spans"):
            asmt_spans = ds.get_spans(field=pids[:])

        for k in mask_keys:
            with utils.Timer("getting per patient mask summary for {}".format(k)):
                writer = ds.get_numeric_writer(asmt_psum, k, 'int8')
                ds.apply_spans_max(asmt_spans, ds.get_reader(f_asmts[k])[:], writer)
                print(sorted(utils.build_histogram(ds.get_reader(asmt_psum[k])[:])))

        for k in isolation_keys:
            with utils.Timer("getting per patient isolation summary for {}".format(k)):
                writer = ds.get_numeric_writer(asmt_psum, k, 'int32')
                ds.apply_spans_max(asmt_spans, ds.get_reader(f_asmts[k])[:], writer)
                print(sorted(utils.build_histogram(ds.get_reader(asmt_psum[k])[:])))

        with utils.Timer("getting prediction maxes for patients"):
            p_predictions = predictions.get_writer(asmt_psum, 'prediction')
            ds.apply_spans_max(asmt_spans, predictions, p_predictions)
            p_predictions = ds.get_reader(asmt_psum[k])
            positives = p_predictions[:] > 0
            print("max covid prediction:", np.count_nonzero(positives), len(positives))

        with utils.Timer("getting patient ids from assessments"):
            writer = pids.get_writer(asmt_psum, 'patient_id')
            writer.write(pd.unique(pids[:]))
    else:
        asmt_psum = tmp['patient_assessment_summaries']

    s_ptnts = src['patients']
    print(s_ptnts.keys())

    pdf = pd.DataFrame({'id': ds.get_reader(s_ptnts['id'])[:],
                        'hwwc': ds.get_reader(s_ptnts['health_worker_with_contact'])[:]})
    adf = pd.DataFrame({'patient_id': ds.get_reader(asmt_psum['patient_id'])[:]})
    jdf = pd.merge(left=adf, right=pdf, left_on='patient_id', right_on='id', how='left')
    print(len(jdf['hwwc']))
    class TestResults:
        def __init__(self):
            self.positive = 0
            self.total = 0

        def add(self, result):
            if result:
                self.positive += 1
            self.total += 1

    results = defaultdict(TestResults)
    positives = ds.get_reader(asmt_psum['prediction'])[:]
    positives = positives > 0
    mask_0 = ds.get_reader(asmt_psum['mask_cloth_or_scarf'])[:]
    mask_1 = ds.get_reader(asmt_psum['mask_surgical'])[:]
    mask_2 = ds.get_reader(asmt_psum['mask_cloth_or_scarf'])[:]
    # mask = mask_0 | mask_1 | mask_2
    mask = mask_0
    print(np.unique(mask, return_counts=True))
    isol_lots = ds.get_reader(asmt_psum['isolation_lots_of_people'])[:]
    isol_lots_7 = np.where(isol_lots > 7, 7, isol_lots)
    print(np.unique(isol_lots_7, return_counts=True))
    print(len(mask), len(positives), len(isol_lots_7))

    # isolation lots of users
    for i_r in range(len(mask)):
        results[(isol_lots_7[i_r], mask[i_r])].add(positives[i_r])

    groupings = sorted(list((r[0], (r[1].positive, r[1].total)) for r in results.items()))

    for g in groupings:
        print(g[0], g[1][0], g[1][1], g[1][0] / g[1][1])

def ppe_use_and_travel_2(ds, src, dest, start_ts):
    ds = session.Session()
    s_ptnts = src['patients']
    s_asmts = src['assessments']
    print(s_asmts.keys())
    s_tests = src['tests']

    if 'filtered_patients' not in dest.keys():
        f_ptnts = dest.create_group('filtered_patients')
        f_asmts = dest.create_group('filtered_assessments')
        f_tests = dest.create_group('filtered_tests')

        # calculate patient first positives
        raw_p_ids = ds.get(s_ptnts['id']).data[:]
        raw_p_acts = ds.get(s_ptnts['assessment_count']).data[:]
        raw_a_pids = ds.get(s_asmts['patient_id']).data[:]
        raw_t_pids = ds.get(s_tests['patient_id']).data[:]

        # filter out anyone without assessments
        patient_filter = raw_p_acts > 0

        print("patient_filter:",
              np.count_nonzero(patient_filter), np.count_nonzero(patient_filter == 0))

        # filter patients
        f_p_ids = ds.get(s_ptnts['id']).create_like(f_ptnts, 'id')
        f_p_ids.data.write(ds.apply_filter(patient_filter, raw_p_ids))

        # filter out any orphaned assessments
        with utils.Timer("fk in pk"):
            assessment_filter = persistence.foreign_key_is_in_primary_key(raw_p_ids, raw_a_pids)
        print("assessment_filter:",
              np.count_nonzero(assessment_filter), np.count_nonzero(assessment_filter == False))
        f_a_pids = ds.get(s_asmts['patient_id']).create_like(f_asmts, 'patient_id')
        f_a_pids.data.write(ds.apply_filter(assessment_filter, raw_a_pids))
        for k in ('created_at', 'tested_covid_positive'):
            field = ds.get(s_asmts[k]).create_like(f_asmts, k)
            field.data.write(ds.apply_filter(assessment_filter, ds.get(s_asmts[k]).data[:]))

        # filter out any orphaned tests
        test_filter = persistence.foreign_key_is_in_primary_key(raw_p_ids, raw_t_pids)
        print("test_filter:",
              np.count_nonzero(test_filter), np.count_nonzero(test_filter == False))
        f_t_pids = ds.get(s_tests['patient_id']).create_like(f_tests, 'patient_id')
        f_t_pids.data.write(ds.apply_filter(test_filter, raw_t_pids))

    else:
        f_ptnts = dest['filtered_patients']
        f_asmts = dest['filtered_assessments']
        f_tests = dest['filtered_tests']
        f_p_ids = ds.get(f_ptnts['id'])
        f_a_pids = ds.get(f_asmts['patient_id'])
        f_t_pids = ds.get(f_tests['patient_id'])

    # calculate the shared set of indices for assessments / tests back to patients
    with utils.Timer("get_shared_index"):
        p_inds, a_pinds, t_pinds = ds.get_shared_index((f_p_ids, f_a_pids, f_t_pids))
    print(max(p_inds.max(), a_pinds.max(), t_pinds.max()))

    # now filter only assessments with positive test results
    pos_asmt_tests = ds.get(f_asmts['tested_covid_positive']).data[:] == 3
    print("old tests positive:",
          np.count_nonzero(pos_asmt_tests), np.count_nonzero(pos_asmt_tests == False))

    # now filter only tests with positive test results

    s_asmts = src['assessments']
    a_cats = ds.get(f_asmts['created_at'])
    asmt_filter = a_cats.data[:] >= start_ts
    print(np.count_nonzero(asmt_filter), len(asmt_filter))
    raw_a_cats = ds.apply_filter(asmt_filter, a_cats.data[:])
    a_days = np.zeros(len(raw_a_cats), dtype=np.int32)
    start_dt = datetime.fromtimestamp(start_ts)
    for i_r in range(len(raw_a_cats)):
        a_days[i_r] = (datetime.fromtimestamp(raw_a_cats[i_r]) - start_dt).days
    print(sorted(utils.build_histogram(a_days)))

if __name__ == '__main__':
    datastore = DataStore()
    src_file = '/home/ben/covid/ds_20200901_full.hdf5'
    dest_file = '/home/ben/covid/ds_20200901_ppe.hdf5'
    with h5py.File(src_file, 'r') as src_data:
        with h5py.File(dest_file, 'w') as dest_data:
            start_timestamp = datetime.timestamp(datetime(2020, 6, 12))
            ppe_use_and_travel_2(datastore, src_data, dest_data, start_timestamp)


#  Imports

import tkinter as tk
from tkinter import *
import cv2
from PIL import Image, ImageTk
import os
import numpy as np
import cv2
from keras.models import Sequential
from keras.layers import Dense, Dropout, Flatten
from keras.layers import Conv2D
from keras.optimizers import Adam
from keras.layers import MaxPooling2D
from keras.preprocessing.image import ImageDataGenerator
import threading

emotion_model = Sequential()

emotion_model.add(Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=(48 ,48 ,1)))
emotion_model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
emotion_model.add(MaxPooling2D(pool_size=(2, 2)))
emotion_model.add(Dropout(0.25))

emotion_model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
emotion_model.add(MaxPooling2D(pool_size=(2, 2)))
emotion_model.add(Conv2D(128, kernel_size=(3, 3), activation='relu'))
emotion_model.add(MaxPooling2D(pool_size=(2, 2)))
emotion_model.add(Dropout(0.25))

emotion_model.add(Flatten())
emotion_model.add(Dense(1024, activation='relu'))
emotion_model.add(Dropout(0.5))
emotion_model.add(Dense(7, activation='softmax'))
emotion_model.load_weights('model.h5')
cv2.ocl.setUseOpenCL(False)

#  GUI and mapping with emojis:

emotion_dict = {0: "Sad", 1: "Disgusted", 2: "Fearful", 3: "Happy", 4: "Neutral", 5: "Angry", 6: "Surprised"}

emoji_dis t ={0 :"C:\\Users\\Saurav\\Desktop\\emojify 2\\emojis\\sad.png"
              ,1 :"C:\\Users\\Saurav\\Desktop\\emojify 2\\emojis\\disgust.png"
              ,2 :"C:\\Users\\Saurav\\Desktop\\emojify 2\\emojis\\fear.png"
              ,3 :"C:\\Users\\Saurav\\Desktop\\emojify 2\\emojis\\happy.png"
              ,4 :"C:\\Users\\Saurav\\Desktop\\emojify 2\\emojis\\neutral.png"
              ,5 :"C:\\Users\\Saurav\\Desktop\\emojify 2\\emojis\\angry.png"
              ,6 :"C:\\Users\\Saurav\\Desktop\\emojify 2\\emojis\\surprised.png"}

global last_frame1
last1_frame1 = np.zeros((448 ,640 ,3), dtype = np.uint8)
global cap1
show_text = [0]
global frame_number

#  Using openCV haarcascade xml to detect the bounding boxes of face in the webcam and predict the emotions:

def show_subject():
    cap1 = cv2.VideoCapture(0)
    if not cap1.isOpened():
        print("Can't open the camera")
    global frame_number
    length = int(cap1.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_number += 1
    if frame_number >= length:
        exit()
    cap1.set(1, frame_number)
    flag1, frame1 = cap1.read()
    frame1 = cv2.resize(frame1 ,(600 ,500))
    bounding_box = cv2.CascadeClassifier \
        ("C:\\Users\\Saurav\\Desktop\\emojify 2\\opencv-master\\opencv-master\\data\\haarcascades\\haarcascade_frontalface_default.xml")
    gray_frame = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    num_faces = bounding_box.detectMultiScale(gray_frame, scaleFactor = 1.3, minNeighbors = 5)
    for (x ,y ,w ,h) in num_faces:
        cv2.rectangle(frame1, (x , y -50), ( x +w, y+ h + 10), (255, 0, 0), 2)
        roi_gray_frame = gray_frame[y:y + h, x:x + w]
        cropped_img = np.expand_dims(np.expand_dims(cv2.resize(roi_gray_frame, (48, 48)), -1), 0)
        prediction = emotion_model.predict(cropped_img)
        maxindex = int(np.argmax(prediction))
        cv2.putText(frame1, emotion_dict[maxindex], (x + 20, y - 60), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2,
                    cv2.LINE_AA)
        show_text[0] = maxindex

    if flag1 is None:
        print("Major error!")
    elif flag1:
        global last_frame1
        last_frame1 = frame1.copy()
        pic = cv2.cvtColor(last_frame1, cv2.COLOR_BGR2RGB)
        img = Image.fromarray(pic)
        imgtk = ImageTk.PhotoImage(image=img)
        lmain.imgtk = imgtk
        lmain.configure(image=imgtk)
        root.update()
        lmain.after(10, show_subject)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        exit()

#  Function to display the emoji after detecting the emotion of the face:

def show_avatar():
    frame2 = cv2.imread(emoji_dist[show_text[0]])
    pic2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2RGB)
    img2 = Image.fromarray(frame2)
    imgtk2 = ImageTk.PhotoImage(image=img2)
    lmain2.imgtk2 = imgtk2
    lmain3.configure(text=emotion_dict[show_text[0]], font=('arial', 45, 'bold'))

    lmain2.configure(image=imgtk2)
    root.update()
    lmain2.after(10, show_avatar())

if __name__ == '__main__':
    frame_number = 0
    root = tk.Tk()
    lmain = tk.Label(master=root, padx=50, bd=10)
    lmain2 = tk.Label(master=root, bd=10)
    lmain3 = tk.Label(master=root, bd=10, fg="#CDCDCD", bg='black')
    lmain.pack(side=LEFT)
    lmain.place(x=50, y=250)
    lmain3.pack()
    lmain3.place(x=960, y=250)
    lmain2.pack(side=RIGHT)
    lmain2.place(x=900, y=350)

    root.title("Photo To Emoji")
    root.geometry("1400x900+100+10")
    root['bg'] = 'black'
    exitButton = Button(root, text='Quit', fg="red", command=root.destroy, font=('arial', 25, 'bold')).pack(side=BOTTOM)
    threading.Thread(target=show_subject()).start()
    threading.Thread(target=show_avatar()).start()
    root.mainloop()


from soccerdepth.data.dataset_loader import get_set
import numpy as np
import utils.files as file_utils
from os.path import join
import argparse
from soccerdepth.models.hourglass import  hg8
from soccerdepth.models.utils import weights_init
from soccerdepth.data.data_utils import image_logger_converter_visdom
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data import DataLoader
import torch.backends.cudnn as cudnn
from soccerdepth.data.transforms import *
from torchvision import transforms
import warnings
from visdom import Visdom

warnings.filterwarnings("ignore")

parser = argparse.ArgumentParser(description='Depth AutoEncoder')
parser.add_argument('--dataset', default='', help='Dataset to train on')
parser.add_argument('--batchSize', type=int, default=6, help='training batch size')
parser.add_argument('--testBatchSize', type=int, default=2, help='testing batch size')
parser.add_argument('--nEpochs', type=int, default=1000, help='number of epochs to train for')
parser.add_argument('--input_nc', type=int, default=4, help='input image channels')
parser.add_argument('--output_nc', type=int, default=51, help='output image channels')
parser.add_argument('--nf', type=int, default=64, help='number of filters')
parser.add_argument('--lr', type=float, default=0.00003, help='Learning Rate. Default=0.002')
parser.add_argument('--beta1', type=float, default=0.9, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda', action='store_true', help='use cuda?')
parser.add_argument('--threads', type=int, default=4, help='number of threads for data loader to use')
parser.add_argument('--seed', type=int, default=123, help='random seed to use. Default=123')
parser.add_argument('--lamb', type=int, default=100, help='weight on L1 term in objective')
parser.add_argument('--run', type=int, default=44, help='Run number for tensorboard')
parser.add_argument('--output_path', default='/home/krematas/Mountpoints/grail/CNN', help='Where the files WILL be stored')
parser.add_argument('--modeldir', default='/home/krematas/Mountpoints/grail/CNN', help='Where the files ARE being stored')
parser.add_argument('--additional_input', default='mask', help='filepostfix')
parser.add_argument('--postfix', default='hg_estmask', help='filepostfix')
parser.add_argument('--resume', type=int, default=0, help='Resume training')
parser.add_argument('--port', type=int, default=9876, help='Run number for tensorboard')
parser.add_argument('--additional_input_type', default='estmask', help='The type of addtional type to load [estmap, trimap]')

opt, _ = parser.parse_known_args()
opt.cuda = True
print(opt)

# Initiate 5 windows
viz = Visdom(env='FIFA CNN training', port=opt.port)
viz.close()
win0 = viz.images(np.ones((4, 3, 128, 128)))
win1 = viz.images(np.ones((4, 3, 128, 128)))
win2 = viz.images(np.ones((4, 3, 128, 128)))
win3 = viz.images(np.ones((4, 3, 128, 128)))
epoch_lot = viz.line(X=torch.zeros((1,)).cpu(), Y=torch.zeros((1,)).cpu(),
                     opts=dict(
                            xlabel='Epoch',
                            ylabel='Loss',
                            title='Epoch Training Loss',
                            legend=['Loss'])
                    )
lot = viz.line(
            X=torch.zeros((1,)).cpu(),
            Y=torch.zeros((1,)).cpu(),
            opts=dict(
                xlabel='Iteration',
                ylabel='Loss',
                title='Current Training Loss',
                legend=['Loss']
            )
        )

# writer = SummaryWriter("runs/run%d" % opt.run)

if opt.cuda and not torch.cuda.is_available():
    raise Exception("No GPU found, please run without --cuda")

cudnn.benchmark = True

torch.manual_seed(opt.seed)
if opt.cuda:
    torch.cuda.manual_seed(opt.seed)

dataset = 'play_for_data'
path_to_data = join(file_utils.get_platform_datadir(dataset), 'cnn2')

print('===> Loading datasets')
composed = transforms.Compose([RandomRotation(), RandomCrop(), Rescale(256, 64), ColorOffset(), ToTensor(), NormalizeImage()])
train_set = get_set(join(path_to_data, 'train'), nbins=opt.output_nc, transform=composed, additional_input_type=opt.additional_input_type)

composed = transforms.Compose([Rescale(256, 64), ToTensor(), NormalizeImage()])
test_set = get_set(join(path_to_data, 'test'), nbins=opt.output_nc, transform=composed, additional_input_type=opt.additional_input_type)

training_data_loader = DataLoader(dataset=train_set, num_workers=8, batch_size=opt.batchSize, shuffle=True)
testing_data_loader = DataLoader(dataset=test_set, num_workers=8, batch_size=opt.testBatchSize, shuffle=False)

print('===> Building model')
model = hg8(input_nc=opt.input_nc, output_nc=opt.output_nc)
model.apply(weights_init)

print('===> The loss function is cross entropy loss')
class_weights = np.ones((opt.output_nc, ))
class_weights[0] = 0.1
weights = torch.from_numpy(class_weights)
weights = torch.FloatTensor(weights.size()).copy_(weights).cuda()

criterion = nn.NLLLoss2d(weight=weights)
logsoftmax = nn.LogSoftmax()

# Setup the Adam optimizer
optimizer = optim.Adam(model.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999), weight_decay=0.003)

# Resume if available
if opt.resume > 0:
    checkpoint = torch.load(join(opt.modeldir, 'model_epoch_%d_%s_%s.pth' % (opt.resume, opt.additional_input, opt.postfix)))
    model.load_state_dict(checkpoint['state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer'])

if opt.cuda:
    model = model.cuda()
    criterion = criterion.cuda()

def train(epoch):

    epoch_loss = 0
    for iteration, batch in enumerate(training_data_loader, 1):
        input, target, mask = Variable(batch['image']).float(), Variable(batch['target']).long(), Variable(batch['mask']).float()
        if opt.input_nc > 3:
            input = torch.cat((input, mask), 1)

        if opt.cuda:
            input = input.cuda()
            target = target.cuda()

        output = model(input)
        loss = criterion(logsoftmax(output[0]), target.squeeze())
        for j in range(1, len(output)):
            loss += criterion(logsoftmax(output[j]), target.squeeze())

        epoch_loss += loss.data[0]
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print("===> Epoch[{}]({}/{}): Loss: {:.4f}".format(epoch, iteration, len(training_data_loader), loss.data[0]))

        if iteration % 50 == 0:
            prediction = logsoftmax(output[-1])

            x, y, z, w = image_logger_converter_visdom(input, mask, target, prediction)

            viz.images(
                x,
                win=win0,
            )
            viz.images(
                y,
                win=win1,
            )
            viz.images(
                w,
                win=win2,
            )
            viz.images(
                z,
                win=win3,
            )

            # writer.add_image('Train_image', x, epoch)
            # writer.add_image('Train_prediction', y, epoch)
            # writer.add_image('Train_label', z, epoch)

        # print(torch.ones((1,)).cpu().size())
        # print(torch.Tensor([loss.data[0]]).unsqueeze(0).cpu().size())
        viz.line(
            X=torch.ones((1, 1)).cpu() * ((epoch-1)*len(training_data_loader) + iteration),
            Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(),
            win=lot,
            update='append'
        )

        # hacky fencepost solution for 0th epoch plot
        if epoch == 1 and iteration == len(training_data_loader)-1:
            viz.line(
                X=torch.zeros((1, 1)).cpu(),
                Y=torch.Tensor([loss.data[0]]).unsqueeze(0).cpu(),
                win=epoch_lot,
                update=True
            )

    viz.line(
        X=torch.ones((1, 1)).cpu()*epoch,
        Y=torch.Tensor([epoch_loss/len(training_data_loader)]).unsqueeze(0).cpu(),
        win=epoch_lot,
        update='append'
    )
    print("===> Epoch {} Complete: Avg. Loss: {:.6f}".format(epoch, epoch_loss / len(training_data_loader)))
    return epoch_loss / len(training_data_loader)

def test():

    epoch_loss = 0
    for iteration, batch in enumerate(testing_data_loader):
        input, target, mask = Variable(batch['image']).float(), Variable(batch['target']).long(), Variable(batch['mask']).float()

        if opt.input_nc > 3:
            input = torch.cat((input, mask), 1)

        if opt.cuda:
            input = input.cuda()
            target = target.cuda()

        output = model(input)

        loss = 0
        for o in output:
            loss += criterion(logsoftmax(o), target.squeeze())

        epoch_loss += loss.data[0]

        if iteration == 1:
            prediction = logsoftmax(output[-1])

            x, y, z, w = image_logger_converter_visdom(input, mask, target, prediction)

            viz.images(
                x,
                win=win0,
            )
            viz.images(
                y,
                win=win1,
            )
            viz.images(
                w,
                win=win2,
            )
            viz.images(
                z,
                win=win3,
            )

    print("===> Validation Complete: Avg. Loss: {:.6f}".format(epoch_loss / len(testing_data_loader)))
    return epoch_loss / len(testing_data_loader)

def checkpoint(epoch):
    model_out_path = "{0}/model_epoch_{1}_{2}_{3}.pth".format(opt.output_path, epoch, opt.additional_input, opt.postfix)
    dict_to_save = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
    }
    torch.save(dict_to_save, model_out_path)
    print("Checkpoint saved to {}".format(model_out_path))

for epoch in range(opt.resume+1, opt.nEpochs + 1):
    v1 = train(epoch)
    v2 = test()
    writer.add_scalar('train_loss', v1, epoch)
    writer.add_scalar('val_loss', v2, epoch)
    checkpoint(epoch)

writer.close()


# -*- coding: utf-8 -*-
"""
Created on Sun Jun  9 17:19:51 2019

@author: Tobias
"""
import statsmodels.sandbox.distributions.extras as gh
def split(a, n):
    k, m = divmod(len(a), n)
    return (a[i * k + min(i, m):(i + 1) * k + min(i + 1, m)] for i in range(n))

os.chdir('C:\\Users\\Tobias\\Dropbox\\Master\\U.S. Data\\Returns')
with open('sp500prices.pickle', 'rb') as file:
       prices=pickle.load(file)

os.chdir('C:\\Users\\Tobias\\Dropbox\\Master\\U.S. Data\\Returns')
indX=pd.read_csv('C:\\Users\\Tobias\\Dropbox\\Master\\U.S. Data\\Returns\\sp500returns.csv')

os.chdir('C:\\Users\\Tobias\\Dropbox\\Master\\U.S. Data\\Similarity\\NewsSim\\40')
os.chdir('C:\\Users\\Tobias\\Dropbox\\Master\\U.S. Data\\Similarity\\40')

with open('ranking.pickle', 'rb') as file:
    tickers = pickle.load(file)

os.chdir('C:\\Users\\Tobias\\Dropbox\\Master\\TempFolder')
with open('ranking25news.pickle','rb') as file:
     tickers = pickle.load(file)

res1   = []
res2   = []
res3   = []
res4   = []
res5   = []
res6   = []
res7   = []
res8   = []
res9   = []
res10  = []
res11   = []
res12   = []
res13   = []
res14   = []
res15  = []
resind = []
resall = []
resff3= []

w1   = []
w2   = []
w3   = []
w4   = []
w5   = []
w6   = []
w7   = []
w8   = []
w9   = []
w10  = []
w11  = []
w12  = []
w13  = []
w14  = []
w15  = []

prices    = pd.read_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\Returns\ReturnData.csv',index_col=[0])#return_data
value     = pd.read_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\Returns\ValueData.csv',index_col=[0])#value_data
indX      = pd.DataFrame()
start     = dt.datetime(1996, 1, 1)
end       = dt.datetime(2019,1,1)
indX['^SP500TR'] = web.get_data_yahoo('^SP500TR',dt.datetime(1996, 1, 1),end=dt.datetime(2019,1,1))['Close']
indX       = np.log(indX/indX.shift(1))
tickers   = tickers.loc['1999-07-22':]#tickers.loc[:]:'2010-02-24'
dates     = tickers.index.to_list()
available = list(prices.columns.values)
valids    = [v for v in list(tickers.columns.values) if v in available]
tickers   = tickers[valids]
val       = None

for ii in range(len(dates)-1):
    data     = tickers.iloc[ii,:]
    data     = data.dropna()
    data     = data.sort_values(axis=0)
    uniques  = data.unique()

    lucky_ten = list(split(uniques,15))
    d1  = []
    d2  = []
    d3  = []
    d4  = []
    d5  = []
    d6  = []
    d7  = []
    d8  = []
    d9  = []
    d10  = []
    d11 = []
    d12 = []
    d13 = []
    d14 = []
    d15 = []
    for idx, m in enumerate(lucky_ten):
        if idx==0:
            for idy in range(len(m)):
                d1.append(data.loc[data==m[idy]].index.values[0])

        if idx==1:
            for idy in range(len(m)):
                d2.append(data.loc[data==m[idy]].index.values[0])

        if idx==2:
            for idy in range(len(m)):
                d3.append(data.loc[data==m[idy]].index.values[0])
        if idx==3:
            for idy in range(len(m)):
                d4.append(data.loc[data==m[idy]].index.values[0])
        if idx==4:
            for idy in range(len(m)):
                d5.append(data.loc[data==m[idy]].index.values[0])
        if idx==5:
            for idy in range(len(m)):
                d6.append(data.loc[data==m[idy]].index.values[0])
        if idx==6:
            for idy in range(len(m)):
                d7.append(data.loc[data==m[idy]].index.values[0])
        if idx==7:
            for idy in range(len(m)):
                d8.append(data.loc[data==m[idy]].index.values[0])
        if idx==8:
             for idy in range(len(m)):
                d9.append(data.loc[data==m[idy]].index.values[0])
        if idx==9:
            for idy in range(len(m)):
                d10.append(data.loc[data==m[idy]].index.values[0])
        if idx==10:
            for idy in range(len(m)):
                d11.append(data.loc[data==m[idy]].index.values[0])
        if idx==11:
            for idy in range(len(m)):
                d12.append(data.loc[data==m[idy]].index.values[0])
        if idx==12:
            for idy in range(len(m)):
                d13.append(data.loc[data==m[idy]].index.values[0])
        if idx==13:
             for idy in range(len(m)):
                d14.append(data.loc[data==m[idy]].index.values[0])
        if idx==14:
            for idy in range(len(m)):
                d15.append(data.loc[data==m[idy]].index.values[0])

    portfolio1_r =    np.sum(np.log(prices[d1].loc[dates[ii]:dates[ii+1]]+1))
    portfolio2_r =    np.sum(np.log(prices[d2].loc[dates[ii]:dates[ii+1]]+1))
    portfolio3_r =    np.sum(np.log(prices[d3].loc[dates[ii]:dates[ii+1]]+1))
    portfolio4_r =    np.sum(np.log(prices[d4].loc[dates[ii]:dates[ii+1]]+1))
    portfolio5_r =    np.sum(np.log(prices[d5].loc[dates[ii]:dates[ii+1]]+1))
    portfolio6_r =    np.sum(np.log(prices[d6].loc[dates[ii]:dates[ii+1]]+1))
    portfolio7_r =    np.sum(np.log(prices[d7].loc[dates[ii]:dates[ii+1]]+1))
    portfolio8_r =    np.sum(np.log(prices[d8].loc[dates[ii]:dates[ii+1]]+1))
    portfolio9_r =    np.sum(np.log(prices[d9].loc[dates[ii]:dates[ii+1]]+1))
    portfolio10_r=    np.sum(np.log(prices[d10].loc[dates[ii]:dates[ii+1]]+1))

    portfolio11_r =    np.sum(np.log(prices[d11].loc[dates[ii]:dates[ii+1]]+1))
    portfolio12_r =    np.sum(np.log(prices[d12].loc[dates[ii]:dates[ii+1]]+1))
    portfolio13_r =    np.sum(np.log(prices[d13].loc[dates[ii]:dates[ii+1]]+1))
    portfolio14_r =    np.sum(np.log(prices[d14].loc[dates[ii]:dates[ii+1]]+1))
    portfolio15_r =    np.sum(np.log(prices[d15].loc[dates[ii]:dates[ii+1]]+1))

    ret_ind      =    np.sum(np.log(indX.loc[dates[ii]:dates[ii+1]]+1))
    ret_all      =    np.sum(np.log(prices.loc[dates[ii]:dates[ii+1]]+1))
#    ret_ff3      =    np.sum(np.log(data3.loc[dates[ii]:dates[ii+1]]/100+1))

    if val:
        portfolio1_v  = value[d1].loc[dates[ii]]/value[d1].loc[dates[ii]].sum()
        portfolio2_v  = value[d2].loc[dates[ii]]/value[d2].loc[dates[ii]].sum()
        portfolio3_v  = value[d3].loc[dates[ii]]/value[d3].loc[dates[ii]].sum()
        portfolio4_v  = value[d4].loc[dates[ii]]/value[d4].loc[dates[ii]].sum()
        portfolio5_v  = value[d5].loc[dates[ii]]/value[d5].loc[dates[ii]].sum()
        portfolio6_v  = value[d6].loc[dates[ii]]/value[d6].loc[dates[ii]].sum()
        portfolio7_v  = value[d7].loc[dates[ii]]/value[d7].loc[dates[ii]].sum()
        portfolio8_v  = value[d8].loc[dates[ii]]/value[d8].loc[dates[ii]].sum()
        portfolio9_v  = value[d9].loc[dates[ii]]/value[d9].loc[dates[ii]].sum()
        portfolio10_v = value[d10].loc[dates[ii]]/value[d10].loc[dates[ii]].sum()
        portfolio11_v = value[d11].loc[dates[ii]]/value[d11].loc[dates[ii]].sum()
        portfolio12_v = value[d12].loc[dates[ii]]/value[d12].loc[dates[ii]].sum()
        portfolio13_v = value[d13].loc[dates[ii]]/value[d13].loc[dates[ii]].sum()
        portfolio14_v = value[d14].loc[dates[ii]]/value[d14].loc[dates[ii]].sum()
        portfolio15_v = value[d15].loc[dates[ii]]/value[d15].loc[dates[ii]].sum()
        if ii==0:
            print('CALCULATING VALUE-WEIGHTED RETURNS')

    else:
        try:
            portfolio1_v  = 1/len(d1)
            portfolio2_v  = 1/len(d2)
            portfolio3_v  = 1/len(d3)
            portfolio4_v  = 1/len(d4)
            portfolio5_v  = 1/len(d5)
            portfolio6_v  = 1/len(d6)
            portfolio7_v  = 1/len(d7)
            portfolio8_v  = 1/len(d8)
            portfolio9_v  = 1/len(d9)
            portfolio10_v = 1/len(d10)
            portfolio11_v = 1/len(d11)
            portfolio12_v = 1/len(d12)
            portfolio13_v = 1/len(d13)
            portfolio14_v = 1/len(d14)
            portfolio15_v = 1/len(d15)
            if ii==0:
                print('CALCULATING EQUAL-WEIGHTED RETURNS')
        except:
            print(ii)
            continue

    w1.append(value[d1].loc[dates[ii]])
    w2.append(value[d2].loc[dates[ii]])
    w3.append(value[d3].loc[dates[ii]])
    w4.append(value[d4].loc[dates[ii]])
    w5.append(value[d5].loc[dates[ii]])
    w6.append(value[d6].loc[dates[ii]])
    w7.append(value[d7].loc[dates[ii]])
    w8.append(value[d8].loc[dates[ii]])
    w9.append(value[d9].loc[dates[ii]])
    w10.append(value[d10].loc[dates[ii]])
    w11.append(value[d11].loc[dates[ii]])
    w12.append(value[d12].loc[dates[ii]])
    w13.append(value[d13].loc[dates[ii]])
    w14.append(value[d14].loc[dates[ii]])
    w15.append(value[d15].loc[dates[ii]])

    ret_1 = np.sum(portfolio1_r*portfolio1_v)
    ret_2 = np.sum(portfolio2_r*portfolio2_v)
    ret_3 = np.sum(portfolio3_r*portfolio3_v)
    ret_4 = np.sum(portfolio4_r*portfolio4_v)
    ret_5 = np.sum(portfolio5_r*portfolio5_v)
    ret_6 = np.sum(portfolio6_r*portfolio6_v)
    ret_7 = np.sum(portfolio7_r*portfolio7_v)
    ret_8 = np.sum(portfolio8_r*portfolio8_v)
    ret_9 = np.sum(portfolio9_r*portfolio9_v)
    ret_10 = np.sum(portfolio10_r*portfolio10_v)
    ret_11 = np.sum(portfolio11_r*portfolio11_v)
    ret_12 = np.sum(portfolio12_r*portfolio12_v)
    ret_13 = np.sum(portfolio13_r*portfolio13_v)
    ret_14 = np.sum(portfolio14_r*portfolio14_v)
    ret_15 = np.sum(portfolio15_r*portfolio15_v)

    res1.append(ret_1)
    res2.append(ret_2)
    res3.append(ret_3)
    res4.append(ret_4)
    res5.append(ret_5)
    res6.append(ret_6)
    res7.append(ret_7)
    res8.append(ret_8)
    res9.append(ret_9)
    res10.append(ret_10)
    res11.append(ret_11)
    res12.append(ret_12)
    res13.append(ret_13)
    res14.append(ret_14)
    res15.append(ret_15)
    resall.append(ret_all)
 #   resff3.append(ret_ff3)
    print(ret_2)
    #res5.append(ret_5)
    #res6.append(ret_6)
    #res7.append(ret_7)
    resind.append(ret_ind)

res1  = pd.DataFrame(res1)
res2  = pd.DataFrame(res2)
res3  = pd.DataFrame(res3)
res4  = pd.DataFrame(res4)
res5  = pd.DataFrame(res5)
res6  = pd.DataFrame(res6)
res7  = pd.DataFrame(res7)
res8  = pd.DataFrame(res8)
res9  = pd.DataFrame(res9)
res10 = pd.DataFrame(res10)
res11 = pd.DataFrame(res11)
res12 = pd.DataFrame(res12)
res13 = pd.DataFrame(res13)
res14 = pd.DataFrame(res14)
res15 = pd.DataFrame(res15)

resall = pd.DataFrame(resall)
ff3    = pd.DataFrame(resff3)
ff3.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF5\ff5.csv')

w1 = pd.DataFrame(w1)
w2 = pd.DataFrame(w2)
w3 = pd.DataFrame(w3)
w4 = pd.DataFrame(w4)
w5 = pd.DataFrame(w5)
w6 = pd.DataFrame(w6)
w7 = pd.DataFrame(w7)
w8 = pd.DataFrame(w8)
w9 = pd.DataFrame(w9)
w10= pd.DataFrame(w10)

[np.mean(w1.mean()),
np.mean(w2.mean()),
np.mean(w3.mean()),
np.mean(w4.mean()),
np.mean(w5.mean()),
np.mean(w6.mean()),
np.mean(w7.mean()),
np.mean(w8.mean()),
np.mean(w9.mean()),
np.mean(w10.mean())]

[len(w1.var()),
 len(w2.var()),
 len(w3.var()),
 len(w4.var()),
 len(w5.var()),
 len(w6.var()),
 len(w7.var()),
 len(w8.var()),
 len(w9.var()),
 len(w10.var())]

#res6 = pd.DataFrame(res6)
#res7 = pd.DataFrame(res7)
resind = pd.DataFrame(resind)
res_ben=res10-res1
res_ben=pd.concat( [pd.Series([0]),res_ben] )
res_ben.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port.csv')
#resind.to_pickle(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\Returns\ReturnSPTRI.pickle')
res1.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port1.csv')
res2.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port2.csv')
res3.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port3.csv')
res4.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port4.csv')
res5.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port5.csv')
res6.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port6.csv')
res7.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port7.csv')
res8.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port8.csv')
res9.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port9.csv')
res10.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\10port10.csv')

res1['2'] = res2
res1['3'] = res3
res1['4'] = res4
res1['5'] = res5
res1['6'] = res6
res1['7'] = res7
res1['8'] = res8
res1['9'] = res9
res1['10'] = res10
resALLport = res1.rename(index=str, columns={0: 1})
resALLport.to_csv(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\FF3\5portall.csv')

###############################################################################

a=res1.mean()
b=res2.mean()
c=res3.mean()
d=res4.mean()
e=res5.mean()
f=res6.mean()
g=res7.mean()
h=res8.mean()
i=res9.mean()
j=res10.mean()

[a,b,c,d,e,f,g,h,i,j]

k=res1.std()
l=res2.std()
m=res3.std()
n=res4.std()
o=res5.std()
p=res6.std()
q=res7.std()
r=res8.std()
s=res9.std()
t=res10.std()

[k,l,m,n,o,p,q,r,s,t]

u=res1.skew()
v=res2.skew()
w=res3.skew()
x=res4.skew()
y=res5.skew()
z=res6.skew()
aa=res7.skew()
bb=res8.skew()
cc=res9.skew()
dd=res10.skew()

[u,v,w,x,y,z,aa,bb,cc,dd]

ee=res1.kurtosis()
ff=res2.kurtosis()
gg=res3.kurtosis()
hh=res4.kurtosis()
ii=res5.kurtosis()
jj=res6.kurtosis()
kk=res7.kurtosis()
ll=res8.kurtosis()
mm=res9.kurtosis()
nn=res10.kurtosis()

[ee,ff,gg,hh,ii,jj,kk,ll,mm,nn]

longshort=res10-res1
longshort.mean()
longshort.std()
longshort.skew()
longshort.kurtosis()
longshort.mean()/longshort.std()

###############################################################################

###############################################################################

res1 = pd.DataFrame(res1)
res2 = pd.DataFrame(res2)
res3 = pd.DataFrame(res3)
res4 = pd.DataFrame(res4)
res5 = pd.DataFrame(res5)

a=res1.mean()
b=res2.mean()
c=res3.mean()
d=res4.mean()
e=res5.mean()

[a,b,c,d,e]

k=res1.std()
l=res2.std()
m=res3.std()
n=res4.std()
o=res5.std()

[k,l,m,n,o]

u=res1.skew()
v=res2.skew()
w=res3.skew()
x=res4.skew()
y=res5.skew()

[u,v,w,x,y]

ee=res1.kurtosis()
ff=res2.kurtosis()
gg=res3.kurtosis()
hh=res4.kurtosis()
ii=res5.kurtosis()

[ee,ff,gg,hh,ii]

longshort=res5-res1
longshort.mean()
longshort.std()
longshort.skew()
longshort.kurtosis()

w1 = pd.DataFrame(w1)
w2 = pd.DataFrame(w2)
w3 = pd.DataFrame(w3)
w4 = pd.DataFrame(w4)
w5 = pd.DataFrame(w5)

[np.mean(w1.mean()),
np.mean(w2.mean()),
np.mean(w3.mean()),
np.mean(w4.mean()),
np.mean(w5.mean())]

[len(w1.var()),
 len(w2.var()),
 len(w3.var()),
 len(w4.var()),
 len(w5.var())]

longshort=res5-res1
longshort.mean()
longshort.std()
longshort.skew()
longshort.kurtosis()
longshort.mean()/longshort.std()

plt.scatter([1,2,3,4,5],[a,b,c,d,e])
plt.xlabel('Portfolio')
plt.ylabel('Return')
plt.title('Value-Weighted Returns')

resind = pd.DataFrame(resind)
resind['ls'] = longshort
print('Sharpe:' +str(longshort.mean()/longshort.std()) + '\n' + 'Corr:' + str(resind.corr()))
###############################################################################

###############################################################################

a=res1.mean()
b=res2.mean()
c=res3.mean()
d=res4.mean()
e=res5.mean()
f=res6.mean()
g=res7.mean()
h=res8.mean()
i=res9.mean()
j=res10.mean()
f1=res11.mean()
g1=res12.mean()
h1=res13.mean()
i1=res14.mean()
j1=res15.mean()

[a,b,c,d,e,f,g,h,i,j,f1,g1,h1,i1,j1]

k=res1.std()
l=res2.std()
m=res3.std()
n=res4.std()
o=res5.std()
p=res6.std()
q=res7.std()
r=res8.std()
s=res9.std()
t=res10.std()
p1=res11.std()
q1=res12.std()
r1=res13.std()
s1=res14.std()
t1=res15.std()

[k,l,m,n,o,p,q,r,s,t,p1,q1,r1,s1,t1]

u=res1.skew()
v=res2.skew()
w=res3.skew()
x=res4.skew()
y=res5.skew()
z=res6.skew()
aa=res7.skew()
bb=res8.skew()
cc=res9.skew()
dd=res10.skew()
z1=res11.skew()
aa1=res12.skew()
bb1=res13.skew()
cc1=res14.skew()
dd1=res15.skew()

[u,v,w,x,y,z,aa,bb,cc,dd,z1,aa1,bb1,cc1,dd1]

ee=res1.kurtosis()
ff=res2.kurtosis()
gg=res3.kurtosis()
hh=res4.kurtosis()
ii=res5.kurtosis()
jj=res6.kurtosis()
kk=res7.kurtosis()
ll=res8.kurtosis()
mm=res9.kurtosis()
nn=res10.kurtosis()
jj1=res11.kurtosis()
kk1=res12.kurtosis()
ll1=res13.kurtosis()
mm1=res14.kurtosis()
nn1=res15.kurtosis()

[ee,ff,gg,hh,ii,jj,kk,ll,mm,nn,jj1,kk1,ll1,mm1,nn1]

longshort=res15-res1
longshort.mean()
longshort.std()
longshort.skew()
longshort.kurtosis()
longshort.mean()/longshort.std()

resind = pd.DataFrame(resind)
resind['ls'] = longshort
print('Sharpe:' +str(longshort.mean()/longshort.std()) + '\n' + 'Corr:' + str(resind.corr()))

w1 = pd.DataFrame(w1)
w2 = pd.DataFrame(w2)
w3 = pd.DataFrame(w3)
w4 = pd.DataFrame(w4)
w5 = pd.DataFrame(w5)
w6 = pd.DataFrame(w6)
w7 = pd.DataFrame(w7)
w8 = pd.DataFrame(w8)
w9 = pd.DataFrame(w9)
w10= pd.DataFrame(w10)
w11 = pd.DataFrame(w11)
w12 = pd.DataFrame(w12)
w13 = pd.DataFrame(w13)
w14 = pd.DataFrame(w14)
w15= pd.DataFrame(w15)

[np.mean(w1.mean()),
np.mean(w2.mean()),
np.mean(w3.mean()),
np.mean(w4.mean()),
np.mean(w5.mean()),
np.mean(w6.mean()),
np.mean(w7.mean()),
np.mean(w8.mean()),
np.mean(w9.mean()),
np.mean(w10.mean()),
np.mean(w11.mean()),
np.mean(w12.mean()),
np.mean(w13.mean()),
np.mean(w14.mean()),
np.mean(w15.mean())]

[len(w1.var()),
 len(w2.var()),
 len(w3.var()),
 len(w4.var()),
 len(w5.var()),
 len(w6.var()),
 len(w7.var()),
 len(w8.var()),
 len(w9.var()),
 len(w10.var()),
 len(w11.var()),
 len(w12.var()),
 len(w13.var()),
 len(w14.var()),
 len(w15.var())]

plt.scatter([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15],[a,b,c,d,e,f,g,h,i,j,f1,g1,h1,i1,j1])
plt.xlabel('Portfolio')
plt.ylabel('Return')
plt.title('Equal-Weighted Returns')

###############################################################################

resind.mean()
resind.std()
resind.skew()
resind.kurtosis()

plt.subplot(210)

# equivalent but more general
ax1=plt.subplot(2, 2,1)

# add a subplot with no frame
ax2=plt.subplot(222, frameon=False)

# add a polar subplot
plt.subplot(223, projection='polar')

# add a red subplot that shares the x-axis with ax1
plt.subplot(224, sharex=ax1, facecolor='red')

# delete ax2 from the figure
plt.delaxes(ax2)

# add ax2 to the figure again
plt.subplot(ax2)

v=j-a

n=np.linspace(-3,3,10000)
sorting = []
indexob = []
pdf_sort = gh.pdf_mvsk( [res7.mean()[0],res7.std()[0]**2,res7.skew()[0],res7.kurt()[0]-3])
pdf_indx = gh.pdf_mvsk( [resind.mean()[0],resind.std()[0]**2,resind.skew()[0],resind.kurt()[0]-3])
for ii in n:
    sorting.append(pdf_sort(ii))
    indexob.append(pdf_indx(ii))

a=res1.std()
b=res2.std()
c=res3.std()
d=res4.std()
e=res5.std()
f=res6.std()
g=res7.std()
h=res8.std()
i=res9.std()
j=res10.std()

plt.scatter([1,2,3,4,5,6,7,8,9,10],[a,b,c,d,e,f,g,h,i,j])
plt.xlabel('Portfolio')
plt.ylabel('Return')
plt.title('Value-Weighted Returns')

res_ben1=pd.concat( [pd.Series([0]),res_ben] )
res_ben2=pd.concat( [pd.Series([0]),res_ben])
res_ind1=pd.concat( [pd.Series([0]),resind] )

with open(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\Returns\ReturnlongshortValueWWithNews.pickle', 'rb') as file:
       ValuWNews = pickle.load(file)
       ValuWNews.index=range(39)

with open(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\Returns\ReturnlongshortEqualWWithNews.pickle', 'rb') as file:
       EqualWNews = pickle.load(file)
       EqualWNews.index=range(39)

with open(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\Returns\ReturnlongshortEqualWWithoutNews.pickle', 'rb') as file:
       EqualW = pickle.load(file)
       EqualW.index=range(39)

with open(r'C:\Users\Tobias\Dropbox\Master\U.S. Data\Returns\ReturnlongshortValueWWithoutNews.pickle', 'rb') as file:
       ValueW = pickle.load(file)
       ValueW.index=range(1,39)

pIndex = np.exp(np.cumsum(resind))
p1     = np.exp(np.cumsum(ValuWNews))
p2     = np.exp(np.cumsum(EqualWNews))
ax = plt.subplot(1,1,1)
plot1, = ax.plot(pIndex, label="line 1")
plot2, = ax.plot(p1 , label="line 2")
plot3, = ax.plot(p2, label="line 3")
ax.legend(['S&P 500 Total Return Index','Value-Weighted Portfolios','Equal-Weighted Portfolios'])
plt.ylabel('Return')
plt.title('FOMC With News')
plt.savefig('FOMCWithNews.png')

p3     = np.exp(np.cumsum(EqualW))
p4     = np.exp(np.cumsum(ValueW))
ax = plt.subplot(1,1,1)
plot1, = ax.plot(pIndex, label="line 1")
plot2, = ax.plot(p3 , label="line 2")
plot3, = ax.plot(p4, label="line 3")
ax.legend(['S&P 500 Total Return Index','Value-Weighted Portfolios','Equal-Weighted Portfolios'])
plt.ylabel('Return')
plt.title('FOMC Without News')
plt.savefig('FOMCWithoutNews.png')

ValuWNews=ValuWNews.iloc[1:]
ValuWNews.index=resind.index
EqualWNews=EqualWNews.iloc[1:]
EqualWNews.index=resind.index
ValueW=ValueW.iloc[1:]
ValueW.index=resind.index
EqualW=EqualW.iloc[1:]
EqualW.index=resind.index
resind['ValueWNews'] = ValuWNews
resind['EqualWNews'] = EqualWNews
resind['ValueNNews'] = ValueW
resind['EqualNNews'] = EqualW

res1 = pd.DataFrame(res1)
res2 = pd.DataFrame(res2)
res3 = pd.DataFrame(res3)
res4 = pd.DataFrame(res4)
res5 = pd.DataFrame(res5)
res6 = pd.DataFrame(res6)
res7 = pd.DataFrame(res7)
res8 = pd.DataFrame(res8)
res9 = pd.DataFrame(res9)
res10= pd.DataFrame(res10)

p = pd.DataFrame({'p1':res1.iloc[:,0]})
p['p2'] = res2
p['p3'] = res3
p['p4'] = res4
p['p5'] = res5
p['p6'] = res6
p['p7'] = res7
p['p8'] = res8
p['p9'] = res9
p['p10'] = res10
p['p11'] = res11
p['p12'] = res12
p['p13'] = res13
p['p14'] = res14
p['p15'] = res15

p=resind
p['wport']=res_ben

ret_mean_an=p.mean()*2
vcov=p.cov()*2

## Identify mean-variance portfolio through simulation ##
risk_free   = web.get_data_fred('TB3MS',start=end,end=end)/100
### Identify GMVP by constrained optimization ###
rf = risk_free.values[0].tolist()
rf = rf[0]
rf=0

# Initial guesses
x0   = np.ones(len(p.var()))
x0  /= x0.sum()

# Run optimization routine, find GMVP
b    = (-1,1)
bnds = (b,)*len(p.var())
con  = {'type': 'eq', 'fun': constraint}
gmvp = minimize(objective_gmvp,x0,args=(ret_mean_an,vcov,rf),constraints=con,
                bounds=bnds,options={'disp': True})
var_gmvp = np.matmul(np.matmul(gmvp.x,vcov),np.transpose(gmvp.x))
ret_gmvp = np.matmul(gmvp.x,ret_mean_an)
std_gmvp = var_gmvp**0.5
s_p_gmvp = (ret_gmvp-rf)/std_gmvp
print('.'*100)
print('\nGlobal Mean Variance portfolio:\n')
print('Return = ' + str(ret_gmvp) + '\n')
print('Standard deviation = ' +  str(std_gmvp) + '\n')
print('Sharpe ratio = ' + str(s_p_gmvp) + '\n')
print('.'*100)

p = pd.DataFrame({'p1':res1.iloc[:,0]})
p['p2'] = res2
p['p3'] = res3
p['p4'] = res4
p['p5'] = res5
p['p6'] = res6
p['p7'] = res7
p['p8'] = res8
p['p9'] = res9
p['p10'] = res10
p['p11'] = res11
p['p12'] = res12
p['p13'] = res13
p['p14'] = res14
p['p15'] = res15
p['spreadport'] = p['p15']-p['p1']

constant_reg = smf.ols(formula=' p1 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p2 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p3 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p4 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p5 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p6 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p7 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p8 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p9 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p10 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p11 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p12 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p13 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p14 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' p15 ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()

constant_reg = smf.ols(formula=' spreadport ~ 1', data=p).fit()
f=constant_reg.get_robustcov_results(cov_type='HAC',maxlags=6)
f.summary()


#! /usr/bin/env python3
#
#  Copyright 2018 California Institute of Technology
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
# ISOFIT: Imaging Spectrometer Optimal FITting
# Author: David R Thompson, david.r.thompson@jpl.nasa.gov
#

from sys import platform
import json
import os
import re
import scipy as s
from common import json_load_ascii, combos, VectorInterpolator
from common import recursive_replace
from copy import deepcopy
from scipy.stats import norm as normal
from scipy.interpolate import interp1d
from rt_lut import TabularRT, FileExistsError

eps = 1e-5  # used for finite difference derivative calculations

class ModtranRT(TabularRT):
    """A model of photon transport including the atmosphere."""

    def __init__(self, config, instrument):

        TabularRT.__init__(self, config, instrument)

        self.modtran_dir = self.find_basedir(config)
        self.filtpath = os.path.join(self.lut_dir, 'wavelengths.flt')
        self.template = deepcopy(json_load_ascii(
            config['modtran_template_file'])['MODTRAN'])

        # Insert aerosol templates, if specified
        if 'aerosol_template_file' in config:
            self.template[0]['MODTRANINPUT']['AEROSOLS'] = \
                deepcopy(json_load_ascii(config['aerosol_template_file']))

        # Insert aerosol data, if specified
        if 'aerosol_model_file' in config:
            aer_data = s.loadtxt(config['aerosol_model_file'])
            self.aer_wl = aer_data[:, 0]
            aer_data = aer_data[:, 1:].T
            self.naer = int(len(aer_data)/3)
            aer_absc, aer_extc, aer_asym = [], [], []
            for i in range(self.naer):
                aer_extc.append(aer_data[i*3])
                aer_absc.append(aer_data[i*3+1])
                aer_asym.append(aer_data[i*3+2])
            self.aer_absc = s.array(aer_absc)
            self.aer_extc = s.array(aer_extc)
            self.aer_asym = s.array(aer_asym)

        # Build the lookup table
        self.build_lut(instrument)

    def find_basedir(self, config):
        '''Seek out a modtran base directory'''

        try:
            return config['modtran_directory']
        except KeyError:
            pass  # fall back to environment variable
        try:
            return os.getenv('MODTRAN_DIR')
        except KeyError:
            raise KeyError('I could not find the MODTRAN base directory')

    def load_tp6(self, infile):
        '''Load a .tp6 file.  This contains the solar geometry.  We
           Return cosine of mean solar zenith'''

        with open(infile, 'r') as f:
            ts, te = -1, -1  # start and end indices
            lines = []
            while len(lines) == 0 or len(lines[-1]) > 0:
                try:
                    lines.append(f.readline())
                except UnicodeDecodeError:
                    pass
            #lines = f.readlines()
            for i, line in enumerate(lines):
                if "SINGLE SCATTER SOLAR" in line:
                    ts = i+5
                if ts >= 0 and len(line) < 5:
                    te = i
                    break
            if ts < 0:
                raise ValueError('Could not find solar geometry in .tp6 file')
        szen = s.array([float(lines[i].split()[3])
                        for i in range(ts, te)]).mean()
        return szen

    def load_chn(self, infile, coszen):
        """Load a .chn output file and parse critical coefficient vectors.
           These are:
             wl      - wavelength vector
             sol_irr - solar irradiance
             sphalb  - spherical sky albedo at surface
             transm  - diffuse and direct irradiance along the
                          sun-ground-sensor path
             transup - transmission along the ground-sensor path only
           We parse them one wavelength at a time."""

        with open(infile) as f:
            sols, transms, sphalbs, wls, rhoatms, transups = [], [], [], [], [], []
            lines = f.readlines()
            for i, line in enumerate(lines):
                if i < 5:
                    continue
                toks = line.strip().split(' ')
                toks = re.findall(r"[\S]+", line.strip())
                wl, wid = float(toks[0]), float(toks[8])  # nm
                solar_irr = float(toks[18]) * 1e6 * \
                    s.pi / wid / coszen  # uW/nm/sr/cm2
                rdnatm = float(toks[4]) * 1e6  # uW/nm/sr/cm2
                rhoatm = rdnatm * s.pi / (solar_irr * coszen)
                sphalb = float(toks[23])
                transm = float(toks[22]) + float(toks[21])
                transup = float(toks[24])
                sols.append(solar_irr)
                transms.append(transm)
                sphalbs.append(sphalb)
                rhoatms.append(rhoatm)
                transups.append(rhoatm)
                wls.append(wl)
        params = [s.array(i) for i in
                  [wls, sols, rhoatms, transms, sphalbs, transups]]
        return tuple(params)

    def ext550_to_vis(self, ext550):
        return s.log(50.0) / (ext550 + 0.01159)

    def modtran_driver(self, overrides):
        """Write a MODTRAN 6.0 input file"""

        param = deepcopy(self.template)

        if hasattr(self, 'aer_absc'):
            fracs = s.zeros((self.naer))

        # Perform overrides
        for key, val in overrides.items():
            recursive_replace(param, key, val)

            if key.startswith('AER'):
                i = int(key.split('_')[-1])
                fracs[i] = val

            elif key == 'EXT550' or key == 'AOT550' or key == 'AOD550':
                # MODTRAN 6.0 convention treats negative visibility as AOT550
                recursive_replace(param, 'VIS', -val)

            elif key == 'FILTNM':
                param[0]['MODTRANINPUT']['SPECTRAL']['FILTNM'] = val

            elif key in ['ITYPE', 'H1ALT', 'IDAY', 'IPARM', 'PARM1',
                         'PARM2', 'GMTIME', 'TRUEAZ', 'OBSZEN']:
                param[0]['MODTRANINPUT']['GEOMETRY'][key] = val

        # For custom aerosols, specify final extinction and absorption
        # MODTRAN 6.0 convention treats negative visibility as AOT550
        if hasattr(self, 'aer_absc'):
            total_aot = fracs.sum()
            recursive_replace(param, 'VIS', -total_aot)
            total_extc = self.aer_extc.T.dot(fracs)
            total_absc = self.aer_absc.T.dot(fracs)
            norm_fracs = fracs/(fracs.sum())
            total_asym = self.aer_asym.T.dot(norm_fracs)

            # Normalize to 550 nm
            total_extc550 = interp1d(self.aer_wl, total_extc)(0.55)
            lvl0 = param[0]['MODTRANINPUT']['AEROSOLS']['IREGSPC'][0]
            lvl0['NARSPC'] = len(self.aer_wl)
            lvl0['VARSPC'] = [float(v) for v in self.aer_wl]
            lvl0['ASYM'] = [float(v) for v in total_asym]
            lvl0['EXTC'] = [float(v) / total_extc550 for v in total_extc]
            lvl0['ABSC'] = [float(v) / total_extc550 for v in total_absc]

        return json.dumps({"MODTRAN": param})

    def build_lut(self, instrument, rebuild=False):
        """ Each LUT is associated with a source directory.  We build a
            lookup table by:
              (1) defining the LUT dimensions, state vector names, and the grid
                  of values;
              (2) running modtran if needed, with each MODTRAN run defining a
                  different point in the LUT; and
              (3) loading the LUTs, one per key atmospheric coefficient vector,
                  into memory as VectorInterpolator objects."""

        # Regenerate MODTRAN input wavelength file
        if not os.path.exists(self.filtpath):
            self.wl2flt(instrument.wl, instrument.fwhm, self.filtpath)

        TabularRT.build_lut(self, instrument, rebuild)

    def rebuild_cmd(self, point, fn):

        vals = dict([(n, v) for n, v in zip(self.lut_names, point)])
        vals['DISALB'] = True
        vals['NAME'] = fn
        vals['FILTNM'] = os.path.normpath(self.filtpath)
        modtran_config_str = self.modtran_driver(dict(vals))

        # Check rebuild conditions: LUT is missing or from a different config
        infilename = 'LUT_'+fn+'.json'
        infilepath = os.path.join(self.lut_dir, infilename)
        outchnname = fn+'.chn'
        outchnpath = os.path.join(self.lut_dir, outchnname)
        if not os.path.exists(infilepath) or\
           not os.path.exists(outchnpath):
            rebuild = True
        else:
            with open(infilepath, 'r') as f:
                current = f.read()
                rebuild = (modtran_config_str.strip() != current.strip())

        if not rebuild:
            raise FileExistsError('File exists')

        # write_config_file
        with open(infilepath, 'w') as f:
            f.write(modtran_config_str)

        # Specify location of the proper MODTRAN 6.0 binary for this OS
        xdir = {'linux': 'linux', 'darwin': 'macos', 'windows': 'windows'}
        cmd = self.modtran_dir+'/bin/'+xdir[platform]+'/mod6c_cons '+infilename
        return cmd

    def load_rt(self, point, fn):
        tp6file = self.lut_dir+'/'+fn+'.tp6'
        solzen = self.load_tp6(tp6file)
        coszen = s.cos(solzen * s.pi / 180.0)
        chnfile = self.lut_dir+'/'+fn+'.chn'
        wl, sol, rhoatm, transm, sphalb, transup = self.load_chn(
            chnfile, coszen)
        return wl, sol, solzen, rhoatm, transm, sphalb, transup

    def wl2flt(self, wls, fwhms, outfile):
        """ helper function to generate Gaussian distributions around the center
            wavelengths """
        I = None
        sigmas = fwhms/2.355
        span = 2.0 * (wls[1]-wls[0])  # nm
        steps = 101

        with open(outfile, 'w') as fout:

            fout.write('Nanometer data for sensor\n')
            for wl, fwhm, sigma in zip(wls, fwhms, sigmas):

                ws = wl + s.linspace(-span, span, steps)
                vs = normal.pdf(ws, wl, sigma)
                vs = vs/vs[int(steps/2)]
                wns = 10000.0/(ws/1000.0)

                fout.write('CENTER:  %6.2f NM   FWHM:  %4.2f NM\n' %
                           (wl, fwhm))

                for w, v, wn in zip(ws, vs, wns):
                    fout.write(' %9.4f %9.7f %9.2f\n' % (w, v, wn))