# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import torch
from torch.autograd import Variable
import warnings
from torch import nn
from collections import OrderedDict
import os
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.linear_model import LogisticRegression
import matplotlib.pyplot as plt
import data as data
from data.BehavioralDataset import BehavioralDataset
from data.BehavioralHmSamples import BehavioralHmSamples
import scipy
from sklearn.preprocessing import MinMaxScaler

warnings.filterwarnings("ignore")


def load_netG(path, isize, nz, nc, ngf, n_extra_layers):
    assert isize % 16 == 0, "isize has to be a multiple of 16"

    cngf, tisize = ngf//2, 4
    while tisize != isize:
        cngf = cngf * 2
        tisize = tisize * 2

    main = nn.Sequential()
    # input is Z, going into a convolution
    main.add_module('initial:{0}-{1}:convt'.format(nz, cngf),
                    nn.ConvTranspose2d(nz, cngf, 4, 1, 0, bias=False))
    main.add_module('initial:{0}:batchnorm'.format(cngf),
                    nn.BatchNorm2d(cngf))
    main.add_module('initial:{0}:relu'.format(cngf),
                    nn.ReLU(True))

    csize, cndf = 4, cngf
    while csize < isize//2:
        main.add_module('pyramid:{0}-{1}:convt'.format(cngf, cngf//2),
                        nn.ConvTranspose2d(cngf, cngf//2, 4, 2, 1, bias=False))
        main.add_module('pyramid:{0}:batchnorm'.format(cngf//2),
                        nn.BatchNorm2d(cngf//2))
        main.add_module('pyramid:{0}:relu'.format(cngf//2),
                        nn.ReLU(True))
        cngf = cngf // 2
        csize = csize * 2

    # Extra layers
    for t in range(n_extra_layers):
        main.add_module('extra-layers-{0}:{1}:conv'.format(t, cngf),
                        nn.Conv2d(cngf, cngf, 3, 1, 1, bias=False))
        main.add_module('extra-layers-{0}:{1}:batchnorm'.format(t, cngf),
                        nn.BatchNorm2d(cngf))
        main.add_module('extra-layers-{0}:{1}:relu'.format(t, cngf),
                        nn.ReLU(True))

    main.add_module('final:{0}-{1}:convt'.format(cngf, nc),
                    nn.ConvTranspose2d(cngf, nc, 4, 2, 1, bias=False))
    main.add_module('final:{0}:tanh'.format(nc),
                    nn.Tanh())
    
    state_dict = torch.load(path, map_location=torch.device('cpu'))

    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[5:] # remove `main.`
        new_state_dict[name] = v
    
    main.load_state_dict(new_state_dict, strict=False)
    
    return main


def load_netG_mlp(path, isize, nz, nc, ngf):
    
    main = nn.Sequential(
        # Z goes into a linear of size: ngf
        nn.Linear(nz, ngf),
        nn.ReLU(True),
        nn.Linear(ngf, ngf),
        nn.ReLU(True),
        nn.Linear(ngf, ngf),
        nn.ReLU(True),
        nn.Linear(ngf, nc * isize * isize),
    )
    
    state_dict = torch.load(path, map_location=torch.device('cpu'))

    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[5:] # remove `main.`
        new_state_dict[name] = v
    
    main.load_state_dict(new_state_dict, strict=False)
    
    return main


# +

def load_netD(path, isize, nc, ndf, n_extra_layers):
    
    assert isize % 16 == 0, "isize has to be a multiple of 16"

    main = nn.Sequential()
    # input is nc x isize x isize
    main.add_module('initial:{0}-{1}:conv'.format(nc, ndf),
                    nn.Conv2d(nc, ndf, 4, 2, 1, bias=False))
    main.add_module('initial:{0}:relu'.format(ndf),
                    nn.LeakyReLU(0.2, inplace=True))
    csize, cndf = isize / 2, ndf

    # Extra layers
    for t in range(n_extra_layers):
        main.add_module('extra-layers-{0}:{1}:conv'.format(t, cndf),
                        nn.Conv2d(cndf, cndf, 3, 1, 1, bias=False))
        main.add_module('extra-layers-{0}:{1}:batchnorm'.format(t, cndf),
                        nn.BatchNorm2d(cndf))
        main.add_module('extra-layers-{0}:{1}:relu'.format(t, cndf),
                        nn.LeakyReLU(0.2, inplace=True))

    while csize > 4:
        in_feat = cndf
        out_feat = cndf * 2
        main.add_module('pyramid:{0}-{1}:conv'.format(in_feat, out_feat),
                        nn.Conv2d(in_feat, out_feat, 4, 2, 1, bias=False))
        main.add_module('pyramid:{0}:batchnorm'.format(out_feat),
                        nn.BatchNorm2d(out_feat))
        main.add_module('pyramid:{0}:relu'.format(out_feat),
                        nn.LeakyReLU(0.2, inplace=True))
        cndf = cndf * 2
        csize = csize / 2

    # state size. K x 4 x 4
    main.add_module('final:{0}-{1}:conv'.format(cndf, 1),
                    nn.Conv2d(cndf, 1, 4, 1, 0, bias=False))

    state_dict = torch.load(path, map_location=torch.device('cpu'))

    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[5:] # remove `module.`
        new_state_dict[name] = v

    main.load_state_dict(new_state_dict, strict=False)
    return main
    


# -

def load_netD_mlp(path, isize, nc, ndf):
    
    main = nn.Sequential(
            # Z goes into a linear of size: ndf
            nn.Linear(nc * isize * isize, ndf),
            nn.ReLU(True),
            nn.Linear(ndf, ndf),
            nn.ReLU(True),
            nn.Linear(ndf, ndf),
            nn.ReLU(True),
            nn.Linear(ndf, 1),
        )
    
    state_dict = torch.load(path, map_location=torch.device('cpu'))
    
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k[5:] # remove `module.`
        new_state_dict[name] = v
    
    main.load_state_dict(new_state_dict, strict=False)
    return main


def reshape_mlp_input(input_sample):
    return input_sample.view(input_sample.size(0), 
                             input_sample.size(1) * input_sample.size(2) * input_sample.size(3))


def sample_wrapper(samples):
    input = torch.FloatTensor(samples.shape[0], 1, 32, 32)
    input.resize_as_(samples).copy_(samples)
    return Variable(input)


# +
# load in all needed data
# load real samples
num_models = 10
real_samples_list = []
for i in range(num_models):
    dataset = BehavioralDataset(isCnnData=True, isScoring=True, auto_number=i, niter=50)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=3, shuffle=True, num_workers=1)
    data_iter = iter(dataloader)
    data = data_iter.next()
    next_samples, _ = data
    real_samples_list.append(next_samples)
    

# load fake samples
h_models_samples_list = []
for i in range(1,6):
    dataset = BehavioralHmSamples(modelNum=i, isCnnData=True, isScoring=True)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1000, shuffle=True, num_workers=1)
    data_iter = iter(dataloader)
    data = data_iter.next()
    next_samples, _ = data
    h_models_samples_list.append(next_samples)
# -

# read in discriminators 
num_nets = 10
dataset = 'behavioral'
netD_list = [load_netD('./loss_curves/netD_{0}_50k_{1}_automated.pth'.format(dataset, i), isize=32, nc=1, ndf=64, n_extra_layers=0) for i in range(num_nets)]

# +
# score hierarchical models
# old way of calculating scores
num_h_models = 5
num_real_samples = real_samples_list[0].shape[0]
scores = np.zeros((num_h_models, num_nets))

for j, netD in enumerate(netD_list):
    
    # score real samples 
    real_samples_scores = netD(sample_wrapper(real_samples)).data.numpy()
    
    for i, fake_samples in enumerate(fake_samples_list):
        
        fake_samples_scores = netD(sample_wrapper(fake_samples)).data.numpy()
        
        # see how many fake samples were scored as more real than each real samples
        num_right = sum([np.sum(fake_samples_scores < real_samples_scores[k]) for k in range(num_real_samples)])
#         print('index: ({0}, {1})'.format(i,j))
#         print(num_right)
#         print(np.sum(fake_samples_scores < real_samples_scores[0]))
#         print(np.sum(fake_samples_scores < real_samples_scores[1]))
#         print(np.sum(fake_samples_scores < real_samples_scores[2]))
#         print(num_right / (1000 * real_samples_scores.shape[0]))
#         print('-----------------------')
        scores[i,j] = num_right / (1000 * num_real_samples)
        

# +
# calculate wins array 
# wins indices: 
# i is the index for the real triple and net number
# j corresponds to the fake sample index in a given fake sample vector
# k corresponds to the hierarchical model 
num_h_models = 5
num_real_samples = len(real_samples_list)
num_real_samples_per_set = real_samples_list[0].shape[0]
num_fake_samples = h_models_samples_list[0].shape[0]
array_list = [np.zeros((num_real_samples, num_fake_samples)) for i in range(num_h_models)]
wins = np.dstack(array_list)

for k, h_model_samples in enumerate(h_models_samples_list):
    
    for i, real_samples in enumerate(real_samples_list):

        netD = netD_list[i]
        real_samples_scores = netD(sample_wrapper(real_samples)).data.numpy()
        fake_samples_scores = netD(sample_wrapper(h_model_samples)).data.numpy()
        
        for j in range(num_fake_samples):
            
            wins_for_sample = 0
            
            for m in range(num_real_samples_per_set):
        
                if fake_samples_scores[j] < real_samples_scores[m]:
                    wins_for_sample += 1
        
            wins[i,j,k] = wins_for_sample 
    
# -

np.save('./data/behavioral_wins.npy', wins)

wins = np.load('./data/behavioral_wins.npy')

import json 
pd.Series(wins.tolist()).to_json('./data/behavioral_wins.json', orient='records')

with open('./data/behavioral_wins2.json', 'w+') as f:
    f.write(json.dumps(wins.tolist()))

wins[0,0:20,0]

wins[1,0:20,1]

wins[0:20,0,0]

# +
import codecs, json 

file_path = "./data/behavioral_wins3.json" ## your path variable
json.dump(wins.tolist(), codecs.open(file_path, 'w', encoding='utf-8'), separators=(',', ':'), sort_keys=True, indent=4) 
# -

# # Visualizations of Wins

wins.shape

means_rs_by_hm = np.mean(wins/3, axis=1)

means_rs_by_hm

model_num = []
mean_correct_for_sample = []
for j in range(means_rs_by_hm.shape[1]):
    for i in range(means_rs_by_hm.shape[0]):
        model_num.append(j+1)
        mean_correct_for_sample.append(means_rs_by_hm[i,j])
correct_df = pd.DataFrame(list(zip(model_num, mean_correct_for_sample)), columns=['Hierarchical Model', 'Average Win Percentage'])

import seaborn as sns
sns.set()

# +
ax = sns.scatterplot(x="Hierarchical Model", y="Average Win Percentage", data=correct_df)
ax.set_title('Average Win Percentage by Hierarchical Model')
ax.set_xticklabels(np.linspace(0.5, 5, 10))

# code to modify xticks taken from here:
# https://stackoverflow.com/questions/38947115/how-to-decrease-the-density-of-x-ticks-in-seaborn
for ind, label in enumerate(ax.get_xticklabels()):
    if (ind + 1) % 2 == 0:  # every 10th label is kept
        label.set_visible(True)
    else:
        label.set_visible(False)
# -

np.median(means_rs_by_hm, axis=0)

np.mean(means_rs_by_hm, axis=0)

np.var(means_rs_by_hm, axis=0)

scipy.stats.sem(means_rs_by_hm, axis=0)

# # Analysis of Gammas from R

# $x_{i,j,k}:=$# of wins for fake sample $j$ from model $k$ over real sample set $i$. 
#
# \begin{align*}
# x_{i,j,k}&\sim Poisson(\theta_{i,k})\\
# \theta_{i,k}&=exp(\mu_{i}+\gamma_{k})
# \end{align*}

gamma = pd.read_csv('./data/behavioral_gamma.csv').to_numpy()
mu = pd.read_csv('./data/behavioral_mu.csv').to_numpy()

gamma

mu

# +
index_to_model = {i:'Model {0}'.format(i+1) for i in range(num_h_models)}

places = {i:{j+1:0 for j in range(num_h_models)} for i in index_to_model.values()}

# +

for i in range(gamma.shape[0]):
    
    ranked_indices = np.argsort(gamma[i,:])[::-1]
    placement = 1
    
    for next_best in ranked_indices:
        
        places[index_to_model[next_best]][placement]+=1
        placement += 1
       
# -

places

model_num = []
gamma_list = []
for j in range(gamma.shape[1]):
    for i in range(gamma.shape[0]):
        model_num.append(j+1)
        gamma_list.append(gamma[i,j])
analysis_df = pd.DataFrame(list(zip(model_num, gamma_list)), columns=['Model Number', 'Gamma'])

ax = sns.scatterplot(x="Model Number", y="Gamma", data=analysis_df)
ax.set_title('Gamma vs Model Number')

np.mean(gamma, axis=0)

np.median(gamma, axis=0)

np.var(gamma, axis=0)

scipy.stats.sem(gamma, axis=0)

means = list(np.mean(scores, axis=1))
samp_se_list = list(scipy.stats.sem(scores, axis=1))
var_list = list(np.var(scores, axis=1))

# Where $\hat{p}$ is the probability of a dog getting shocked.
#
# $X_{1it}$ and $X_{2it}$ are the number of times a dog has respectively been shocked and avoided being shocked in previous trials.
#
# Model 1 
# \begin{align*}
# \hat{p}=\sigma(\beta_{1}+\beta_{2}X_{1it}+\beta_{3}X_{2it})
# \end{align*}
#
# Model 2 
# \begin{align*}
# \hat{p}=exp(\beta_{1}X_{1it}+\beta_{2}X_{2it})
# \end{align*}
#
# Model 3 
# \begin{align*}
# \hat{p}&=\sigma(\frac{\alpha}{t}+\gamma)&\text{$t$ is the trial number}
# \end{align*}
#
# Model 4 is the "switch." 
#
# Model 5 
# \begin{align*}
# \hat{p}&=\frac{\alpha}{t}&\text{$t$ is the trial number}
# \end{align*}
#

for mean, var, se, i in zip(means, var_list, samp_se_list, range(1,len(means)+1)):
    print('Model\t {0} \nMean\t {1}\nSE\t {2}\nVar\t {3}'.format(i, mean, se, var))
    print('------------------')


# # Scaled Scores

scaled_scores = scores.copy()
scaler = MinMaxScaler()
scaled_scores = scaler.fit_transform(scaled_scores)
scaled_scores_df = pd.DataFrame(scaled_scores)
scaled_scores_df.to_csv('./data/behavioral_model_scaled_scores.csv', index=False, float_format='%.3f')

scaled_means = list(np.mean(scaled_scores, axis=1))
scaled_samp_se_list = list(scipy.stats.sem(scaled_scores, axis=1))
scaled_var_list = list(np.var(scaled_scores, axis=1))

for mean, var, se, i in zip(scaled_means, scaled_var_list, scaled_samp_se_list, range(1,len(scaled_means)+1)):
    print('Model\t {0} \nMean\t {1}\nSE\t {2}\nVar\t {3}'.format(i, mean, se, var))
    print('------------------')

a = np.zeros((5,1000))
b = np.zeros((5,1000))
np.dstack((a,b)).shape

fake_samples_list[0].shape[0]

num_real_samples

np.sum(a, axis=1).shape

np.reshape(a, (a.shape[1], a.shape[0])).shape

type(a)

a.shape[0]



StarcoderdataJupyterScriptsDedupFiltered