import os
import numpy as np
from os.path import isfile
import torch
import torch.nn.functional as F
EPS = 1e-6

def assert_same_shape(t1, t2):
    for (x, y) in zip(list(t1.shape), list(t2.shape)):
        assert(x==y)

def print_stats_py(name, tensor):
    print('%s (%s) min = %.2f, mean = %.2f, max = %.2f' % (name, tensor.dtype, np.min(tensor), np.mean(tensor), np.max(tensor)))

def tensor2summ(tensor, permute_dim=False):
    # if permute_dim = True:
    # for 2D tensor, assume input is torch format B x S x C x H x W, we want B x S x H x W x C
    # for 3D tensor, assume input is torch format B x S x C x H x W x D, we want B x S x H x W x C x D
    # and finally unbind the sequeence dimension and return a list of [B x H x W x C].
    assert(tensor.ndim == 5 or tensor.ndim == 6)
    assert(tensor.size()[1] == 2) #sequense length should be 2
    if permute_dim:
        if tensor.ndim == 6: #3D tensor
            tensor = tensor.permute(0, 1, 3, 4, 5, 2)
        elif tensor.ndim == 5: #2D tensor
            tensor = tensor.permute(0, 1, 3, 4, 2)

    tensor = torch.unbind(tensor, dim=1)
    return tensor

def normalize_single(d):
    # d is a whatever shape torch tensor
    dmin = torch.min(d)
    dmax = torch.max(d)
    d = (d-dmin)/(EPS+(dmax-dmin))
    return d

def normalize(d):
    # d is B x whatever. normalize within each element of the batch
    out = torch.zeros(d.size())
    if d.is_cuda:
        out = out.cuda()
    B = list(d.size())[0]
    for b in list(range(B)):
        out[b] = normalize_single(d[b])
    return out

def reduce_masked_mean(x, mask, dim=None, keepdim=False):
    # x and mask are the same shape
    # returns shape-1
    # axis can be a list of axes
    assert(x.size() == mask.size())
    prod = x*mask
    if dim is None:
        numer = torch.sum(prod)
        denom = EPS+torch.sum(mask)

    else:
        numer = torch.sum(prod, dim=dim, keepdim=keepdim)
        denom = EPS+torch.sum(mask, dim=dim, keepdim=keepdim)

    mean = numer/denom
    return mean

def pack_seqdim(tensor, B):
    shapelist = list(tensor.shape)
    B_, S = shapelist[:2]
    assert(B==B_)
    otherdims = shapelist[2:]
    tensor = torch.reshape(tensor, [B*S]+otherdims)
    return tensor

def unpack_seqdim(tensor, B):
    shapelist = list(tensor.shape)
    BS = shapelist[0]
    assert(BS%B==0)
    otherdims = shapelist[1:]
    S = int(BS/B)
    tensor = torch.reshape(tensor, [B,S]+otherdims)
    return tensor

def gridcloud3D(B, Z, Y, X, norm=False):
    # we want to sample for each location in the grid
    grid_z, grid_y, grid_x = meshgrid3D(B, Z, Y, X, norm=norm)
    x = torch.reshape(grid_x, [B, -1])
    y = torch.reshape(grid_y, [B, -1])
    z = torch.reshape(grid_z, [B, -1])
    # these are B x N
    xyz = torch.stack([x, y, z], dim=2)
    # this is B x N x 3
    return xyz

def gridcloud2D(B, Y, X, norm=False):
    # we want to sample for each location in the grid
    grid_y, grid_x = meshgrid2D(B, Y, X, norm=norm)
    x = torch.reshape(grid_x, [B, -1])
    y = torch.reshape(grid_y, [B, -1])
    # these are B x N
    xy = torch.stack([x, y], dim=2)
    # this is B x N x 2
    return xy

def gridcloud3D_py(Z, Y, X):
    # we want to sample for each location in the grid
    grid_z, grid_y, grid_x = meshgrid3D_py(Z, Y, X)
    x = np.reshape(grid_x, [-1])
    y = np.reshape(grid_y, [-1])
    z = np.reshape(grid_z, [-1])
    # these are N
    xyz = np.stack([x, y, z], axis=1)
    # this is N x 3
    return xyz

def meshgrid2D_py(Y, X):
    grid_y = np.linspace(0.0, Y-1, Y)
    grid_y = np.reshape(grid_y, [Y, 1])
    grid_y = np.tile(grid_y, [1, X])

    grid_x = np.linspace(0.0, X-1, X)
    grid_x = np.reshape(grid_x, [1, X])
    grid_x = np.tile(grid_x, [Y, 1])

    return grid_y, grid_x

def gridcloud2D_py(Y, X):
    # we want to sample for each location in the grid
    grid_y, grid_x = meshgrid2D_py(Y, X)
    x = np.reshape(grid_x, [-1])
    y = np.reshape(grid_y, [-1])
    # these are N
    xy = np.stack([x, y], axis=1)
    # this is N x 2
    return xy

def normalize_grid3D(grid_z, grid_y, grid_x, Z, Y, X, clamp_extreme=True):
    # make things in [-1,1]
    grid_z = 2.0*(grid_z / float(Z-1)) - 1.0
    grid_y = 2.0*(grid_y / float(Y-1)) - 1.0
    grid_x = 2.0*(grid_x / float(X-1)) - 1.0

    if clamp_extreme:
        grid_z = torch.clamp(grid_z, min=-2.0, max=2.0)
        grid_y = torch.clamp(grid_y, min=-2.0, max=2.0)
        grid_x = torch.clamp(grid_x, min=-2.0, max=2.0)

    return grid_z, grid_y, grid_x

def normalize_grid2D(grid_y, grid_x, Y, X, clamp_extreme=True):
    # make things in [-1,1]
    grid_y = 2.0*(grid_y / float(Y-1)) - 1.0
    grid_x = 2.0*(grid_x / float(X-1)) - 1.0

    if clamp_extreme:
        grid_y = torch.clamp(grid_y, min=-2.0, max=2.0)
        grid_x = torch.clamp(grid_x, min=-2.0, max=2.0)

    return grid_y, grid_x

def normalize_gridcloud3D(xyz, Z, Y, X, clamp_extreme=True):
    # make things in [-1,1]
    x = xyz[...,0]
    y = xyz[...,1]
    z = xyz[...,2]

    z = 2.0*(z / float(Z-1)) - 1.0
    y = 2.0*(y / float(Y-1)) - 1.0
    x = 2.0*(x / float(X-1)) - 1.0

    xyz = torch.stack([x,y,z], dim=-1)

    if clamp_extreme:
        xyz = torch.clamp(xyz, min=-2.0, max=2.0)
    return xyz

def normalize_gridcloud2D(xy, Y, X, clamp_extreme=True):
    # make things in [-1,1]
    x = xy[...,0]
    y = xy[...,1]

    y = 2.0*(y / float(Y-1)) - 1.0
    x = 2.0*(x / float(X-1)) - 1.0

    xy = torch.stack([x,y], dim=-1)

    if clamp_extreme:
        xy = torch.clamp(xy, min=-2.0, max=2.0)
    return xy

def meshgrid3D_yxz(B, Y, X, Z):
    # returns a meshgrid sized B x Y x X x Z
    # this ordering makes sense since usually Y=height, X=width, Z=depth

	grid_y = torch.linspace(0.0, Y-1, Y)
	grid_y = torch.reshape(grid_y, [1, Y, 1, 1])
	grid_y = grid_y.repeat(B, 1, X, Z)

	grid_x = torch.linspace(0.0, X-1, X)
	grid_x = torch.reshape(grid_x, [1, 1, X, 1])
	grid_x = grid_x.repeat(B, Y, 1, Z)

	grid_z = torch.linspace(0.0, Z-1, Z)
	grid_z = torch.reshape(grid_z, [1, 1, 1, Z])
	grid_z = grid_z.repeat(B, Y, X, 1)

	return grid_y, grid_x, grid_z

def meshgrid2D(B, Y, X, stack=False, norm=False):
    # returns a meshgrid sized B x Y x X

    grid_y = torch.linspace(0.0, Y-1, Y, device=torch.device('cuda'))
    grid_y = torch.reshape(grid_y, [1, Y, 1])
    grid_y = grid_y.repeat(B, 1, X)

    grid_x = torch.linspace(0.0, X-1, X, device=torch.device('cuda'))
    grid_x = torch.reshape(grid_x, [1, 1, X])
    grid_x = grid_x.repeat(B, Y, 1)

    if norm:
        grid_y, grid_x = normalize_grid2D(
            grid_y, grid_x, Y, X)

    if stack:
        # note we stack in xy order
        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
        grid = torch.stack([grid_x, grid_y], dim=-1)
        return grid
    else:
        return grid_y, grid_x

def meshgrid3D(B, Z, Y, X, stack=False, norm=False):
    # returns a meshgrid sized B x Z x Y x X

    grid_z = torch.linspace(0.0, Z-1, Z, device=torch.device('cuda'))
    grid_z = torch.reshape(grid_z, [1, Z, 1, 1])
    grid_z = grid_z.repeat(B, 1, Y, X)

    grid_y = torch.linspace(0.0, Y-1, Y, device=torch.device('cuda'))
    grid_y = torch.reshape(grid_y, [1, 1, Y, 1])
    grid_y = grid_y.repeat(B, Z, 1, X)

    grid_x = torch.linspace(0.0, X-1, X, device=torch.device('cuda'))
    grid_x = torch.reshape(grid_x, [1, 1, 1, X])
    grid_x = grid_x.repeat(B, Z, Y, 1)

    if norm:
        grid_z, grid_y, grid_x = normalize_grid3D(
            grid_z, grid_y, grid_x, Z, Y, X)

    if stack:
        # note we stack in xyz order
        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
        grid = torch.stack([grid_x, grid_y, grid_z], dim=-1)
        return grid
    else:
        return grid_z, grid_y, grid_x

def meshgrid3D_py(Z, Y, X, stack=False, norm=False):
    grid_z = np.linspace(0.0, Z-1, Z)
    grid_z = np.reshape(grid_z, [Z, 1, 1])
    grid_z = np.tile(grid_z, [1, Y, X])

    grid_y = np.linspace(0.0, Y-1, Y)
    grid_y = np.reshape(grid_y, [1, Y, 1])
    grid_y = np.tile(grid_y, [Z, 1, X])

    grid_x = np.linspace(0.0, X-1, X)
    grid_x = np.reshape(grid_x, [1, 1, X])
    grid_x = np.tile(grid_x, [Z, Y, 1])

    if norm:
        grid_z, grid_y, grid_x = normalize_grid3D(
            grid_z, grid_y, grid_x, Z, Y, X)

    if stack:
        # note we stack in xyz order
        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
        grid = np.stack([grid_x, grid_y, grid_z], dim=-1)
        return grid
    else:
        return grid_z, grid_y, grid_x

def sub2ind(height, width, y, x):
    return y*width + x

def sql2_on_axis(x, axis, keepdim=True):
    return torch.sum(x**2, axis, keepdim=keepdim)

def l2_on_axis(x, axis, keepdim=True):
    return torch.sqrt(EPS + sql2_on_axis(x, axis, keepdim=keepdim))

def l1_on_axis(x, axis, keepdim=True):
    return torch.sum(torch.abs(x), axis, keepdim=keepdim)

def sub2ind3D(depth, height, width, d, h, w):
    # when gathering/scattering with these inds, the tensor should be Z x Y x X
    return d*height*width + h*width + w

def gradient3D(x, absolute=False, square=False):
    # x should be B x C x D x H x W
    dz = x[:, :, 1:, :, :] - x[:, :, :-1, :, :]
    dy = x[:, :, :, 1:, :] - x[:, :, :, :-1, :]
    dx = x[:, :, :, :, 1:] - x[:, :, :, :, :-1]

    # zeros = tf.zeros_like(x)
    zeros = torch.zeros_like(x)
    zero_z = zeros[:, :, 0:1, :, :]
    zero_y = zeros[:, :, :, 0:1, :]
    zero_x = zeros[:, :, :, :, 0:1]
    dz = torch.cat([dz, zero_z], axis=2)
    dy = torch.cat([dy, zero_y], axis=3)
    dx = torch.cat([dx, zero_x], axis=4)
    if absolute:
        dz = torch.abs(dz)
        dy = torch.abs(dy)
        dx = torch.abs(dx)
    if square:
        dz = dz ** 2
        dy = dy ** 2
        dx = dx ** 2
    return dz, dy, dx

def gradient2D(x, absolute=False, square=False):
    # x should be B x C x H x W
    dh = x[:, :, 1:, :] - x[:, :, :-1, :]
    dw = x[:, :, :, 1:] - x[:, :, :, :-1]

    # zeros = tf.zeros_like(x)
    zeros = torch.zeros_like(x)
    zero_h = zeros[:, :, 0:1, :]
    zero_w = zeros[:, :, :, 0:1]
    dh = torch.cat([dh, zero_h], axis=2)
    dw = torch.cat([dw, zero_w], axis=3)
    if absolute:
        dh = torch.abs(dh)
        dw = torch.abs(dw)
    if square:
        dh = dh ** 2
        dw = dw ** 2
    return dh, dw

def matmul2(mat1, mat2):
    return torch.matmul(mat1, mat2)

def matmul3(mat1, mat2, mat3):
    return torch.matmul(mat1, torch.matmul(mat2, mat3))

def matmul4(mat1, mat2, mat3, mat4):
    return torch.matmul(torch.matmul(mat1, torch.matmul(mat2, mat3)), mat4)

def downsample(img, factor):
    down = torch.nn.AvgPool2d(factor)
    img = down(img)
    return img

def downsample3D(vox, factor):
    down = torch.nn.AvgPool3d(factor)
    vox = down(vox)
    return vox

def downsample3Dflow(flow, factor):
    down = torch.nn.AvgPool3d(factor)
    flow = down(flow) * 1./factor
    return flow

def l2_normalize(x, dim=1):
    # dim1 is the channel dim
    return F.normalize(x, p=2, dim=dim)

def hard_argmax3D(tensor):
    B, Z, Y, X = list(tensor.shape)

    flat_tensor = tensor.reshape(B, -1)
    argmax = torch.argmax(flat_tensor, dim=1)

    # convert the indices into 3D coordinates
    argmax_z = argmax // (Y*X)
    argmax_y = (argmax % (Y*X)) // X
    argmax_x = (argmax % (Y*X)) % X

    argmax_z = argmax_z.reshape(B)
    argmax_y = argmax_y.reshape(B)
    argmax_x = argmax_x.reshape(B)
    return argmax_z, argmax_y, argmax_x

def argmax3D(heat, hard=False):
    B, Z, Y, X = list(heat.shape)

    if hard:
        # hard argmax
        loc_z, loc_y, loc_x = hard_argmax3D(heat)
        loc_z = loc_z.float()
        loc_y = loc_y.float()
        loc_x = loc_x.float()
    else:
        heat = heat.reshape(B, Z*Y*X)
        prob = torch.nn.functional.softmax(heat, dim=1)

        grid_z, grid_y, grid_x = meshgrid3D(B, Z, Y, X)

        grid_z = grid_z.reshape(B, -1)
        grid_y = grid_y.reshape(B, -1)
        grid_x = grid_x.reshape(B, -1)

        loc_z = torch.sum(grid_z*prob, dim=1)
        loc_y = torch.sum(grid_y*prob, dim=1)
        loc_x = torch.sum(grid_x*prob, dim=1)
        # these are B
    return loc_z, loc_y, loc_x


'''
Table of Contents

Functions and Interdependencies:
    proj
    orthogonalize
        - proj
    OLS
    EV
    pairwise_similarity
    best_permutation
        - pairwise_similarity
    self_similarity_pairwise
        - best_permutation
'''
import numpy as np
import scipy.optimize
# import sklearn.decomposition
from numba import njit, prange, jit
import torch
from tqdm import tqdm

from . import indexing

import copy
import time
from functools import partial

def proj(v1, v2):
    '''
    Projects one or more vectors (columns of v1) onto one or more vectors (v2)
    RH 2021

    Args:
        v1 (ndarray):
            vector set 1. Either a single vector or a 2-D array where the columns
             are the vectors
        v2 (ndarray):
            vector set 2. Either a single vector or a 2-D array where the columns
             are the vectors
            If None, v2 is set to v1

    Returns:
        proj_vec (ndarray):
            vector set 1 projected onto vector set 2.
            shape: (v1.shape[0], v1.shape[1], v2.shape[1])
        proj_score (ndarray or scalar):
            projection scores. shape: (v1.shape[1], v2.shape[1])
    '''
    if isinstance(v1, np.ndarray):
        from opt_einsum import contract
        einsum = contract
        norm = partial(np.linalg.norm, axis=0, keepdims=True)
    elif isinstance(v1, torch.Tensor):
        from torch import einsum
        norm = partial(torch.linalg.norm, dim=0, keepdim=True)

    if v2 is None:
        v2 = v1
    if v1.ndim < 2:
        v1 = v1[:,None]
    if v2.ndim < 2:
        v2 = v2[:,None]

    assert v1.shape[0] == v2.shape[0], f"v1.shape[0] must equal v2.shape[0]. Got v1.shape[0]={v1.shape[0]} and v2.shape[0]={v2.shape[0]}"

    u = v2 / norm(v2)
    proj_score = v1.T @ u
    proj_vec = einsum('ik,jk->ijk', u, proj_score)

    return proj_vec, proj_score

def vector_angle(v1, v2=None, mode='cosine_similarity'):
    '''
    Calculates the angle between two vector sets.
    RH 2023

    Args:
        v1 (ndarray):
            vector set 1. Either a single vector or a 2-D array where the columns
             are the vectors
        v2 (ndarray):
            vector set 2. Either a single vector or a 2-D array where the columns
             are the vectors
            If None, v2 is set to v1
        mode (str):
            'deg': returns angle in degrees
            'rad': returns angle in radians
            'cosine_similarity': returns cosine similarity

    Returns:
        angle (ndarray or scalar):
            angle between v1 and v2. shape: (v1.shape[1], v2.shape[1])
    '''
    if isinstance(v1, np.ndarray):
        norm = partial(np.linalg.norm, axis=0, keepdims=True)
        arccos = np.arccos
        rad2deg = np.rad2deg
    elif isinstance(v1, torch.Tensor):
        norm = partial(torch.linalg.norm, dim=0, keepdim=True)
        arccos = torch.acos
        rad2deg = torch.rad2deg

    if v2 is None:
        v2 = v1
    if v1.ndim < 2:
        v1 = v1[:,None]
    if v2.ndim < 2:
        v2 = v2[:,None]

    assert v1.shape[0] == v2.shape[0], f"v1.shape[0] must equal v2.shape[0]. Got v1.shape[0]={v1.shape[0]} and v2.shape[0]={v2.shape[0]}"

    u1 = v1 / norm(v1)
    u2 = v2 / norm(v2)

    if mode == 'cosine_similarity':
        return u1.T @ u2
    elif mode == 'deg':
        return rad2deg(arccos(u1.T @ u2))
    elif mode == 'rad':
        return arccos(u1.T @ u2)
    else:
        raise ValueError(f"mode must be 'cosine_similarity', 'deg', or 'rad'. Got {mode}")

def orthogonalize(v1, v2, method='OLS', device='cpu', thresh_EVR_PCA=1e-15):
    '''
    Orthogonalizes one or more vectors (columns of v1) relative to another set
     of vectors (v2). Subtracts the projection of v1 onto v2 off of v1.
    RH 2021-2023

    Args:
        v1 (ndarray):
            vector set 1. (y_true)
            Either a single vector or a 2-D array where the columns
             are the vectors. shape: (n_samples, n_vectors)
        v2 (ndarray):
            vector set 2. (y_pred)
            Either a single vector or a 2-D array where the columns
             are the vectors. shape: (n_samples, n_vectors)
        method (str):
            'serial': uses a gram-schmidt like iterative process to orthogonalize
             v1 off of v2. This method may have some minor numerical instability
             issues when v2 contains many (>100) vectors.
            'OLS': uses OLS regression to orthogonalize v1 off of v2. This method
             is numerically stable when v2 has many columns, and usually faster.
             However, OLS can have issues when v2 is not full rank or singular.
        device (str):
            Device to use for torch tensors.
            Default is 'cpu'.
        thresh_EVR_PCA (scalar):
            Threshold for the Explained Variance Ratio (EVR) of the PCA components.
            Set so that it might act as a cutoff for redundant components.
            If the EVR of a component is below this threshold, it is not used to
             orthogonalize v1. This is useful when v2 is singular or nearly singular.

    Returns:
        v1_orth (ndarray):
            vector set 1 with the projections onto vector set 2 subtracted off.
             Same size as v1.
        EVR (ndarray):
            Explained Variance Ratio for each column of v1.
            Amount of variance that all the vectors in v2 can explain for each
            vector in v1.
            When v1 is z-scored, EVR is equivalent to pearsons R^2;
            as in pairwise_similarity(OLS(v2, v1)[1] , v1)**2
        EVR_total (scalar):
            total amount of variance explained in v1 by v2
        pca_dict (dict):
            dictionary containing the PCA outputs:
                'comps': PCA components
                'scores': PCA scores
                'singVals': PCA singular values
                'EVR': PCA Explained Variance Ratio
                'PCs_above_thresh': boolean array indicating which PCs are above
                 the threshold for EVR
    '''
    if isinstance(v1, np.ndarray):
        v1 = torch.as_tensor(v1)
        return_numpy = True
    else:
        return_numpy = False
    if isinstance(v2, np.ndarray):
        v2 = torch.as_tensor(v2)

    v1 = v1.to(device)
    v2 = v2.to(device)

    if v1.ndim < 2:
        v1 = v1[:,None]
    if v2.ndim < 2:
        v2 = v2[:,None]

    # I'm pretty sure using PCA is fine for this.
    if v2.shape[1] > 1:
        from .decomposition import torch_pca
        comps, pc_scores, singVals, pc_EVR = torch_pca(
            X_in=v2,
            rank=v2.shape[1],
            mean_sub=True,
            device=device,
            return_cpu=False,
            return_numpy=isinstance(v1, np.ndarray),
            cuda_empty_cache=False
        )
        pca_dict = {
            'comps': comps,
            'scores': pc_scores,
            'singVals': singVals,
            'EVR': pc_EVR,
            'PCs_above_thresh': pc_EVR > thresh_EVR_PCA,
        }
        pc_scores_aboveThreshold = pcsat = pc_scores[:, pca_dict['PCs_above_thresh']]
    else:
        pcsat = v2
        pca_dict = None

    # decomp = sklearn.decomposition.PCA(n_components=v2.shape[1])
    # pc_scores = decomp.fit_transform(v2)

    v1_means = torch.mean(v1, dim=0)
    if method == 'serial':
        # Serial orthogonalization.
        v1_orth = copy.deepcopy(v1 - v1_means)
        for ii in range(pcsat.shape[1]):
            proj_vec = proj(v1_orth , pcsat[:,ii])[0]
            v1_orth = v1_orth.squeeze() - proj_vec.squeeze()
    elif method == 'OLS':
        # Ordinary Least Squares.
        X = torch.cat((pcsat, torch.ones((pcsat.shape[0], 1), dtype=pcsat.dtype, device=device)), dim=1)
        theta = torch.linalg.inv(X.T @ X) @ X.T @ (v1 - v1_means)
        y_rec = X @ theta
        v1_orth = v1 - y_rec + X[:,-1][:,None] * theta[-1]  ## this adds back the bias term

    v1_orth = v1_orth + v1_means

    EVR = 1 - (torch.var(v1_orth, dim=0) / torch.var(v1, dim=0))
    EVR_total = 1 - ( torch.sum(torch.var(v1_orth, dim=0), dim=0) / torch.sum(torch.var(v1, dim=0), dim=0) )

    if return_numpy:
        v1_orth = v1_orth.cpu().numpy()
        EVR = EVR.cpu().numpy()
        EVR_total = EVR_total.cpu().numpy()

        if pca_dict is not None:
            for key in pca_dict.keys():
                if isinstance(pca_dict[key], torch.Tensor):
                    pca_dict[key] = pca_dict[key].cpu().numpy()

    return v1_orth, EVR, EVR_total, pca_dict

@njit
def pair_orth_helper(v1, v2):
    """
    Helper function for main pairwise_orthogonalization
     function. Performs the pairwise orthogonalization
     by subtracting off v2 from v1. Uses numba to speed
     up the computation.
     v1 = v1 - proj(v1 onto v2)
    RH 2021

    Args:
        v1 (ndarray):
            Vector set 1. Columns are vectors.
        v2 (ndarray):
            Vector set 2. Columns are vectors.

    Returns:
        v1_orth (ndarray):
            v1 - proj(v1 onto v2)
    """
    return  v1 - (np.diag(np.dot(v1.T, v2)) / np.diag(np.dot(v2.T, v2))) * v2
def pairwise_orthogonalization(v1, v2, center:bool=True):
    """
    Orthogonalizes columns of v2 off of the columns of v1
     and returns the orthogonalized v1 and the explained
     variance ratio of v2 off of v1.
    v1: y_true, v2: y_pred
    Since it's just pairwise, there should not be any
     numerical instability issues.
    RH 2021

    Args:
        v1 (ndarray):
            y_true
            Vector set 1. Either a single vector or a 2-D
             array where the columns are the vectors.
        v2 (ndarray):
            y_pred
            Vector set 2. Either a single vector or a 2-D
             array where the columns are the vectors.
        center (bool):
            Whether to center the vectors.
            Centering prevents negative EVR values.

    Returns:
        v1_orth (ndarray):
            Vector set 1 with the projections onto vector
             set 2 subtracted off.
            Same size as v1.
        EVR (ndarray):
            Explained Variance Ratio for each column of v1.
            Amount of variance that all the vectors in v2
             can explain for each vector in v1.
        EVR_total_weighted (scalar):
            Average amount of variance explained in v1 by v2
             weighted by the variance of each column of v1.
        EVR_total_unweighted (scalar):
            Average amount of variance explained in v1 by v2
    """
    assert v1.ndim == v2.ndim
    if v1.ndim==1:
        v1 = v1[:,None]
        v2 = v2[:,None]
    assert v1.shape[1] == v2.shape[1]
    assert v1.shape[0] == v2.shape[0]

    if center:
        v1 = v1 - np.mean(v1, axis=0)
        v2 = v2 - np.mean(v2, axis=0)

    v1_orth = pair_orth_helper(v1, v2)

    v1_var = np.var(v1, axis=0)
    EVR = 1 - (np.var(v1_orth, axis=0) / v1_var)

    EVR_total_weighted = np.sum(v1_var * EVR) / np.sum(v1_var)
    EVR_total_unweighted = np.mean(EVR)
    return v1_orth, EVR, EVR_total_weighted, EVR_total_unweighted

@torch.jit.script
def pairwise_orthogonalization_torch_helper(v1, v2, center:bool=True):
    # assert v1.ndim == v2.ndim
    if v1.ndim==1:
        v1 = v1[:,None]
    if v2.ndim==1:
        v2 = v2[:,None]
    # assert v1.shape[1] == v2.shape[1]
    assert v1.shape[0] == v2.shape[0]

    if center:
        v1 = v1 - torch.nanmean(v1, dim=0)
        v2 = v2 - torch.nanmean(v2, dim=0)

    # v1_orth = v1 - (torch.diag(torch.matmul(v1.T, v2)) / torch.diag(torch.matmul(v2.T, v2)))*v2
    v1_orth = v1 - (torch.nansum(v1 * v2, dim=0) / torch.nansum(v2 * v2, dim=0) )*v2

    v1_var = torch.var(v1, dim=0)
    EVR = 1 - (torch.var(v1_orth, dim=0) / v1_var)

    EVR_total_weighted = torch.nansum(v1_var * EVR) / torch.sum(v1_var)
    EVR_total_unweighted = torch.nanmean(EVR)
    return v1_orth, EVR, EVR_total_weighted, EVR_total_unweighted
def pairwise_orthogonalization_torch(v1, v2, center:bool=True, device='cpu'):
    """
    Orthogonalizes columns of v2 off of the columns of v1
     and returns the orthogonalized v1 and the explained
     variance ratio of v2 off of v1.
    Use center=True to get guarantee zero correlation between
     columns of v1_orth and v2.
    v1: y_true, v2: y_pred
    Since it's just pairwise, there should not be any
     numerical instability issues.
    Same as pairwise_orthogonalization, but with
     torch.jit.script instead of njit.
    RH 2021

    Args:
        v1 (ndarray):
            y_true
            Vector set 1. Either a single vector or a 2-D
             array where the columns are the vectors.
        v2 (ndarray):
            y_pred
            Vector set 2. Either a single vector or a 2-D
             array where the columns are the vectors.
        center (bool):
            Whether to center the vectors.
            Centering prevents negative EVR values.

    Returns:
        v1_orth (ndarray):
            Vector set 1 with the projections onto vector
             set 2 subtracted off.
            Same size as v1.
        EVR (ndarray):
            Explained Variance Ratio for each column of v1.
            Amount of variance that all the vectors in v2
             can explain for each vector in v1.
        EVR_total_weighted (scalar):
            Average amount of variance explained in v1 by v2
             weighted by the variance of each column of v1.
        EVR_total_unweighted (scalar):
            Average amount of variance explained in v1 by v2
    """
    if isinstance(v1, np.ndarray):
        v1 = torch.from_numpy(v1)
        return_numpy = True
    else:
        return_numpy = False
    if isinstance(v2, np.ndarray):
        v2 = torch.from_numpy(v2)
    v1 = v1.to(device)
    v2 = v2.to(device)
    v1_orth, EVR, EVR_total_weighted, EVR_total_unweighted = pairwise_orthogonalization_torch_helper(v1, v2, center=center)
    if return_numpy:
        v1_orth = v1_orth.cpu().numpy()
        EVR = EVR.cpu().numpy()
        EVR_total_weighted = EVR_total_weighted.cpu().numpy()
        EVR_total_unweighted = EVR_total_unweighted.cpu().numpy()
    return v1_orth, EVR, EVR_total_weighted, EVR_total_unweighted

def EV(y_true, y_pred):
    '''
    Explained Variance
    Calculating as 1 - (SS_residual_variance / SS_original_variance)
    Should be exactly equivalent to sklearn's
    sklearn.metrics.explained_variance_score but slightly faster and provides
    most relevant outputs all together
    RH 2021

    Args:
        y_true (ndarray):
            array with same size as y_pred. The columns are
            individual y vectors
        y_pred (ndarray):
            array with same size as y_true. The columns are the predicted
            y vectors

    Returns:
        EV (1-D array):
            explained variance of each y_true column by each corresponding
            y_pred column. Same as
            sklearn.metrics.explained_variance_score(y_true, y_pred, multioutput='raw_values')
        EV_total_weighted (scalar):
            total weighted explained variance. Same as
            sklearn.metrics.explained_variance_score(y_true, y_pred, multioutput='variance_weighted')
        EV_total_unweighted (scalar):
            average of all EV values. Same as
            sklearn.metrics.explained_variance_score(y_true, y_pred, multioutput='uniform_average')
    '''

    EV = 1 - np.sum((y_true - y_pred)**2, axis=0) / np.sum((y_true - np.mean(y_true, axis=0))**2, axis=0)
    y_true_var = np.var(y_true, axis=0)
    EV_total_weighted = np.sum( y_true_var* EV ) / np.sum(y_true_var)
    EV_total_unweighted = np.mean(EV)

    return EV , EV_total_weighted , EV_total_unweighted

def pairwise_similarity(v1 , v2=None , method='pearson' , ddof=1):
    '''
    Computes similarity matrices between two sets of vectors (columns within
    2-D arrays) using either covariance, Pearson correlation or cosine_similarity.

    Think of this function as a more general version of np.cov or np.corrcoef
    RH 2021

    Args:
        v1 (ndarray):
            2-D array of column vectors.
        v2 (ndarray):
            2-D array of column vectors to compare to vector_set1. If None, then
             the function is a type of autosimilarity matrix
        method (str):
            'cov' - covariance
            'pearson' or 'R' - Pearson correlation
            'cosine_similarity' - cosine similarity
        ddof (scalar/int):
            Used if method=='cov'. Define the degrees of freedom. Set to 1 for
            unbiased calculation, 0 for biased calculation.
    Returns:
        ouput (ndarray):
            similarity matrix dependent on method
    '''

    methods = ['cov', 'pearson', 'R', 'cosine_similarity']
    assert np.isin(method, methods), f'RH Error: method must be one of: {methods}'

    if v2 is None:
        v2 = v1

    if v1.ndim == 1:
        v1 = v1[:,None]
    if v2.ndim == 1:
        v2 = v2[:,None]

    if method=='cov':
        v1_ms = v1 - np.mean(v1, axis=0) # 'mean subtracted'
        v2_ms = v2 - np.mean(v2, axis=0)
        output = (v1_ms.T @ v2_ms) / (v1.shape[0] - ddof)
    if method in ['pearson', 'R']:
        # Below method should be as fast as numpy.corrcoef .
        # Output should be same within precision, but
        # numpy.corrcoef makes a doublewide matrix and can be annoying

        # Note: Pearson's R-value can be different than sqrt(EV)
        # calculated below if the residuals are not orthogonal to the
        # prediction. Best to use EV for R^2 if unsure
        v1_ms = v1 - np.mean(v1, axis=0) # 'mean subtracted'
        v2_ms = v2 - np.mean(v2, axis=0)
        output = (v1_ms.T @ v2_ms) / np.sqrt(np.sum(v1_ms**2, axis=0, keepdims=True).T * np.sum(v2_ms**2, axis=0, keepdims=True))
    if method=='cosine_similarity':
        output = (v1 / (np.linalg.norm(v1 , axis=0, keepdims=True))).T  @ (v2  / np.linalg.norm(v2 , axis=0, keepdims=True))
    return output

def batched_covariance(X, batch_size=1000, device='cpu'):
    """
    Batched covariance matrix calculation.
    Allows for large datasets to be processed in batches on GPU.
    RH 2022

    Args:
        X (np.ndarray or torch.Tensor):
            2D array of shape (n_samples, n_features)
        batch_size (int):
            Number of samples to process at a time.
        device (str):
            Device to use for computation.

    Returns:
        cov (np.ndarray):
    """
    X_dl1 = list(indexing.make_batches(np.arange(X.shape[1]), batch_size=batch_size, return_idx=True))
    X_dl2 = list(indexing.make_batches(np.arange(X.shape[1]), batch_size=batch_size, return_idx=True))

    if torch.is_tensor(X):
        X_cov = torch.zeros(X.shape[1], X.shape[1], device=device)
    else:
        X_cov = np.zeros((X.shape[1], X.shape[1]))

    n_batches = X.shape[1] // batch_size
    for ii, (X_batch_i, idx_batch_i) in enumerate(tqdm(X_dl1, total=n_batches, leave=False, desc='outer loop')):
        for jj, (X_batch_j, idx_batch_j) in enumerate(tqdm(X_dl2, total=n_batches, leave=False, desc='inner loop')):
            x_t = X[:,idx_batch_i[0]:idx_batch_i[-1]].T
            x   = X[:,idx_batch_j[0]:idx_batch_j[-1]]
            if device != 'cpu':
                x_t = x_t.to(device)
                x = x.to(device)

            X_cov[idx_batch_i[0]:idx_batch_i[-1], idx_batch_j[0]:idx_batch_j[-1]] = x_t @ x
    return X_cov

def batched_matrix_multiply(X1, X2, batch_size1=1000, batch_size2=1000, device='cpu'):
    """
    Batched matrix multiplication of two matrices.
    Allows for multiplying huge matrices together on a GPU.
    RH 2022

    Args:
        X1 (np.ndarray or torch.Tensor):
            first matrix. shape (n_samples1, n_features1).
        X2 (np.ndarray or torch.Tensor):
            second matrix. shape (n_samples2, n_features2).
        batch_size1 (int):
            batch size for first matrix.
        batch_size2 (int):
            batch size for second matrix
        device (str):
            device to use for computation and output

    Returns:
        X1_X2 (np.ndarray or torch.Tensor):
    """
    X1_dl = indexing.make_batches(np.arange(X1.shape[1]), batch_size=batch_size1, return_idx=True)
    X2_dl = indexing.make_batches(np.arange(X2.shape[1]), batch_size=batch_size2, return_idx=True)

    if torch.is_tensor(X1):
        Y = torch.zeros(X1.shape[1], X2.shape[1], device=device)
    else:
        Y = np.zeros((X1.shape[1], X2.shape[1]))

    n_batches1 = X1.shape[1] // batch_size1
    n_batches2 = X2.shape[1] // batch_size2
    for ii, (X_batch_i, idx_batch_i) in enumerate(tqdm(X1_dl, total=n_batches1, leave=False, desc='outer loop')):
        for jj, (X_batch_j, idx_batch_j) in enumerate(X2_dl):
            x1_t = X1[:,idx_batch_i[0]:idx_batch_i[-1]].T
            x2   = X2[:,idx_batch_j[0]:idx_batch_j[-1]]
            if device != 'cpu':
                x1_t = x1_t.to(device)
                x2 = x2.to(device)

            Y[idx_batch_i[0]:idx_batch_i[-1], idx_batch_j[0]:idx_batch_j[-1]] = x1_t @ x2
    return Y
    # X1_dl = indexing.make_batches(X1.T, batch_size=batch_size1, return_idx=True)
    # X2_dl = indexing.make_batches(X2.T, batch_size=batch_size2, return_idx=True)

    # if torch.is_tensor(X1):
    #     Y = torch.zeros(X1.shape[1], X2.shape[1], device=device)
    # else:
    #     Y = np.zeros((X1.shape[1], X2.shape[1]))

    # n_batches1 = X1.shape[1] // batch_size1
    # n_batches2 = X2.shape[1] // batch_size2
    # for ii, (X_batch_i, idx_batch_i) in enumerate(tqdm(X1_dl, total=n_batches1, leave=False, desc='outer loop')):
    #     for jj, (X_batch_j, idx_batch_j) in enumerate(X2_dl):
    #         Y[idx_batch_i[0]:idx_batch_i[-1], idx_batch_j[0]:idx_batch_j[-1]] = X_batch_i @ X_batch_j.T
    # return Y

def similarity_to_distance(x, fn_toUse=1, a=1, b=0, eps=0):
    """
    Convert similarity metric to distance metric.
    RH 2022

    Args:
        x (value or array):
            similarity metric.
        fn_toUse (int from 1 to 7):
            Sets the function to use.
        a (float):
            Scaling parameter.
        b (float):
            Shifting parameter.
        eps (float):
            Small value to add to output.

    """
    if fn_toUse == 1:
        d = ((b+1) / (x**a)) -1 # fn 1: 1/x
    if fn_toUse == 2:
        d = np.exp(((-x+b)**a)) # fn 2: exp(-x)
    if fn_toUse == 3:
        d = np.arctan(a*(-x+b)) # fn 3: arctan(-x)
    if fn_toUse == 4:
        d = b - x**a # fn 4: 1-x
    if fn_toUse == 5:
        d = np.sqrt(1-(x+b))**a # fn 5: sqrt(1-x)
    if fn_toUse == 6:
        d = -np.log((x*a)+b) # fn 6: -log(x)

    return d + eps

def cp_reconstruction_EV(tensor_dense, tensor_CP):
    """
    Explained variance of a reconstructed tensor using
    by a CP tensor (similar to kruskal tensor).
    RH 2023

    Args:
        tensor_dense (np.ndarray or torch.Tensor):
            Dense tensor to be reconstructed. shape (n_samples, n_features)
        tensor_CP (tensorly CPTensor or list of np.ndarray/torch.Tensor):
            CP tensor.
            If a list of factors, then each factor should be a 2D array of shape
            (n_samples, rank).
            Can also be a tensorly CPTensor object.
    """
    tensor_rec = None
    try:
        import tensorly as tl
        if isinstance(tensor_CP, tl.cp_tensor.CPTensor):
            tensor_rec = tl.cp_to_tensor(tensor_CP)
    except ImportError as e:
        raise ImportError('tensorly not installed. Please install tensorly or provide a list of factors as the tensor_CP argument.')

    if tensor_rec is None:
        assert isinstance(tensor_CP, list), 'tensor_CP must be a list of factors'
        assert all([isinstance(f, (np.ndarray, torch.Tensor)) for f in tensor_CP]), 'tensor_CP must be a list of factors'
        tensor_rec = indexing.cp_to_dense(tensor_CP)

    if isinstance(tensor_dense, torch.Tensor):
        var = torch.var
    elif isinstance(tensor_dense, np.ndarray):
        var = np.var
    ev = 1 - (var(tensor_dense - tensor_rec) / var(tensor_dense))
    return ev

##########################################
########### Linear Assignment ############
##########################################

def best_permutation(mat1 , mat2 , method='pearson'):
    '''
    This function compares the representations of two sets of vectors (columns
     of mat1 and columns of mat2).
    We assume that the vectors in mat1 and mat2 are similar up to a permutation.
    We therefore find the 'best' permutation that maximizes the similarity
     between the sets of vectors
    RH 2021

    Args:
        mat1 (np.ndarray):
            a 2D array where the columns are vectors we wish to match with mat2
        mat2 (np.ndarray):
            a 2D array where the columns are vectors we wish to match with mat1
        method (string)  :
            defines method of calculating pairwise similarity between vectors:
                'pearson' or 'cosine_similarity'

    Returns:
        sim_avg (double):
            the average similarity between matched vectors. Units depend on
            method
        sim_matched (double):
            the similarity between each pair of matched vectors.
        ind1 (int):
            indices of vectors in mat1 matched to ind2 in mat2 (usually just
            sequential for ind1)
        ind2 (int):
            indices of vectors in mat2 matched to ind1 in mat1
    '''
    corr = mat1.T @ mat2
    ind1 , ind2 = scipy.optimize.linear_sum_assignment(corr, maximize=True)
    sim_matched = np.zeros(len(ind1))
    for ii in range(len(ind1)):
        if method=='pearson':
            sim_matched[ii] = np.corrcoef(mat1[:,ind1[ii]] , mat2[:,ind2[ii]])[0][1]
        if method=='cosine_similarity':
            sim_matched[ii] = pairwise_similarity( mat1[:,ind1[ii]] , mat2[:,ind2[ii]] , 'cosine_similarity')

    sim_avg = np.mean(sim_matched)
    return sim_avg , sim_matched , ind1.astype('int64') , ind2.astype('int64')

def self_similarity_pairwise(mat_set , method='pearson'):
    '''
    This function compares sets of 2-D matrices within a 3-D array using the
    'best_permutation' function.
    We assume that the vectors within the matrices are similar up to a
    permutation.
    We therefore find the 'best' permutation that maximizes the similarity
    between the sets of vectors within each matrix.
    RH 2021

    Args:
        mat_set (np.ndarray):
            a 3D array where the columns within the first two dims are vectors
             we wish to match with the columns from matrices from other slices
             in the third dimension
        method (string):
            defines method of calculating pairwise similarity between vectors:
             'pearson' or 'cosine_similarity'

    Returns:
        same as 'best_permutation', but over each combo
        combos: combinations of pairwise comparisons
    '''

    import itertools

    n_repeats = mat_set.shape[2]
    n_components = mat_set.shape[1]

    combos = np.array(list(itertools.combinations(np.arange(n_repeats),2)))
    n_combos = len(combos)

    corr_avg = np.zeros((n_combos))
    corr_matched = np.zeros((n_components , n_combos))
    ind1 = np.zeros((n_components , n_combos), dtype='int64')
    ind2 = np.zeros((n_components , n_combos), dtype='int64')
    for i_combo , combo in enumerate(combos):
        corr_avg[i_combo] , corr_matched[:,i_combo] , ind1[:,i_combo] , ind2[:,i_combo]  =  best_permutation(mat_set[:,:,combo[0]]  ,  mat_set[:,:,combo[1]] , method)
    # print(corr_avg)
    return corr_avg, corr_matched, ind1, ind2, combos

def enumerate_paths(edges):
    """
    From Caleb Weinreb 2022

    Create a list of all paths in a directed graph

    Parameters
    ----------
    edges: list of tuples
        Edges in the graph as tuples (i,j) for each edge i->j. The
        edges are assumed to be in topopological order, meaning
        edge (i,j) is listed prior to edge (k,l) whenever j==k.

    Returns
    -------
    paths: list of tuples
        All directed paths as tuples of node indexes
    """
    child_map = {parent:[] for parent,child in edges}
    for parent,child in edges: child_map[parent].append(child)
    leaf_nodes = [child for _,child in edges if not child in child_map]
    sub_paths = {leaf:[[leaf]] for leaf in leaf_nodes}
    for parent,_ in edges[::-1]:
        if not parent in sub_paths:
            sub_paths[parent] = [[parent]]
            for child in child_map[parent]:
                for path in sub_paths[child]:
                    sub_paths[parent].append([parent]+path)
    return [tuple(p) for p in sum(sub_paths.values(),[])]
def maximum_directed_matching(weights, partition):
    """
    From Caleb Weinreb 2022

    Find the "maximum directed matching" in a weighted n-partite graph.

    Let $G$ be a directed graph with nodes $\{1,...,N\}$ and weighted
    edges $E_{ij}$. Assume that the nodes of $G$ are partitioned into
    groups, where $P_i$ denotes the group-label for node $i$, and that
    all the edges $E_{ij}$ satisfy $P_i < P_j$. In other words, all
    edges go from lower to higher partition index.

    We define a collection of node-sets as a "directed matching" if
    each set forms a connected component in $G$ and no node belongs to
    more than one set. The maximum directed matching is defined as the
    directed matching that maximizes the sum edge weights across all
    edges that originate and terminate in the same node-set.

    Parameters
    ----------
    weights : sparse matrix, shape=(N,N)
        The set of edge weights in the graph where N is the total
        number of nodes. ``weights[i,j]`` is the weight of the edge
        (i->j). A weight of 0 implies no edge.

    partition : ndarray, shape=(N,)
        The partition label for each node in the graph. We require
        ``partition[i] < partition[j]`` whenever ``weights[i,j] != 0``

    Returns
    -------
    matching: list of tuples
        Maximum directed matching as a list of node-index tuples.
    """
    import pulp

    # check that the edges are compatible with the partition
    sources,targets = weights.nonzero()
    assert np.all(partition[sources] < partition[targets]), 'Invalid weights/partition combination'

    # enumerate all possible node groupings and their score
    edge_order = np.argsort(partition[sources])
    sorted_edges = list(zip(sources[edge_order], targets[edge_order]))
    paths = enumerate_paths(sorted_edges)
    path_weights = [sum([weights[i,j] for i in p for j in p]) for p in paths]

    # configure an ILP solver with pulp
    problem = pulp.LpProblem('matching', pulp.LpMaximize)

    # initialize variables
    path_vars = pulp.LpVariable.dicts('paths', paths, cat=pulp.LpBinary)
    problem += pulp.lpSum([w*path_vars[p] for p,w in zip(paths, path_weights)])

    # set constraints
    path_membership = {n:[] for n in range(weights.shape[0])}
    for path in paths:
        for n in path: path_membership[n].append(path)
    for n in path_membership:
        problem += pulp.lpSum(path_vars[path] for path in path_membership[n]) <= 1

    # solve and extract results
    problem.solve()
    matching = [p for p,v in path_vars.items() if v.value()==1]
    return matching


import logging
from typing import Type, Tuple, Dict, List
import itertools
from collections import OrderedDict

import torch

import dace
from dace import data as dt

from daceml.autodiff.backward_pass_generator import BackwardPassGenerator
from daceml.autodiff.base_abc import AutoDiffException, BackwardResult
from daceml.onnx.converters import clean_onnx_name
from daceml.onnx.onnx_importer import create_output_array, ONNXModel
from daceml import transformation

log = logging.getLogger(__name__)

def make_backward_function(
    model: ONNXModel,
    required_grads: List[str],
) -> Tuple[dace.SDFG, dace.SDFG, BackwardResult, Dict[str, dt.Data]]:
    """ Convert an ONNXModel to a PyTorch differentiable function. This method should not be used on its own.
        Instead use the ``backward=True`` parameter of :class:`daceml.torch.DaceModule`.

        :param model: the model to convert.
        :param required_grads: the list of inputs names of the module that we must compute gradients for.
        :return: A 4-tuple of forward SDFG, backward SDFG, backward result, and input arrays for
                 backward pass (as mapping of names to DaCe data descriptors).
    """

    if len(model.sdfg.nodes()) != 1:
        raise AutoDiffException(
            "Expected to find exactly one SDFGState, found {}".format(
                len(model.sdfg.nodes())))

    forward_sdfg = model.sdfg
    forward_state = model.sdfg.nodes()[0]

    backward_sdfg = dace.SDFG(forward_sdfg.name + "_backward")
    backward_state = backward_sdfg.add_state()

    gen = BackwardPassGenerator(
        sdfg=forward_sdfg,
        state=forward_state,
        given_gradients=[clean_onnx_name(name) for name in model.outputs],
        required_gradients=required_grads,
        backward_sdfg=backward_sdfg,
        backward_state=backward_state,
        zero_non_transients=False)

    backward_result, backward_grad_arrays, backward_input_arrays = gen.backward(
    )

    replaced_scalars = {}
    for name, desc in backward_input_arrays.items():
        if name not in forward_sdfg.arrays:
            raise AutoDiffException(
                "Expected to find array with name '{}' in SDFG".format(name))

        forward_desc = forward_sdfg.arrays[name]
        # we will save this output and pass it to the backward pass

        # Views should not be forwarded. Instead the backward pass generator should forward the source of the view,
        # and rebuild the sequence of required views in the backward pass.
        assert type(forward_desc) is not dt.View
        if isinstance(forward_desc, dt.Scalar):
            # we can't return scalars from SDFGs, so we add a copy to an array of size 1
            fwd_arr_name, _ = forward_sdfg.add_array(
                name + "_array", [1],
                forward_desc.dtype,
                transient=False,
                storage=forward_desc.storage,
                find_new_name=True)
            bwd_arr_name, bwd_desc = backward_sdfg.add_array(
                name + "_array", [1],
                forward_desc.dtype,
                transient=False,
                storage=forward_desc.storage,
                find_new_name=True)
            backward_sdfg.arrays[name].transient = True

            fwd_copy_state = forward_sdfg.add_state_after(forward_state,
                                                          label="copy_out_" +
                                                          fwd_arr_name)
            bwd_copy_state = backward_sdfg.add_state_before(backward_state,
                                                            label="copy_in_" +
                                                            bwd_arr_name)
            fwd_copy_state.add_edge(fwd_copy_state.add_read(name), None,
                                    fwd_copy_state.add_write(fwd_arr_name),
                                    None, dace.Memlet(name + "[0]"))

            bwd_copy_state.add_edge(bwd_copy_state.add_read(bwd_arr_name),
                                    None, bwd_copy_state.add_write(name), None,
                                    dace.Memlet(name + "[0]"))
            replaced_scalars[name] = (bwd_arr_name, bwd_desc)
        else:
            forward_sdfg.arrays[name].transient = False

    for orig_name, (replaced_name, replaced_desc) in replaced_scalars.items():
        del backward_input_arrays[orig_name]
        backward_input_arrays[replaced_name] = replaced_desc

    for fwd_name, bwd_name in backward_result.required_grad_names.items():
        desc = backward_sdfg.arrays[bwd_name]
        if isinstance(desc, dt.Scalar):
            arr_name, arr_desc = backward_sdfg.add_array(bwd_name + "_array",
                                                         [1],
                                                         desc.dtype,
                                                         transient=False,
                                                         storage=desc.storage,
                                                         find_new_name=True)
            desc.transient = True
            bwd_copy_state = backward_sdfg.add_state_after(backward_state,
                                                           label="copy_out_" +
                                                           bwd_name)
            bwd_copy_state.add_edge(bwd_copy_state.add_read(bwd_name), None,
                                    bwd_copy_state.add_write(arr_name), None,
                                    dace.Memlet(bwd_name + "[0]"))
            backward_result.required_grad_names[fwd_name] = arr_name

    backward_sdfg.validate()

    return forward_sdfg, backward_sdfg, backward_result, backward_input_arrays


from typing import List

def add_border(picture: List[str]) -> List[str]:

    longest = len(max(picture, key=len))

    unframed = [
        '* ' + s + f"{' ' * (longest - (len(s)-1))}*"
        if len(s) < longest
        else f"* {s} *"
        for s in picture
    ]

    edge = '*' * (longest + 4)
    pad = f"*{' ' * (longest + 2)}*"

    return [edge, pad] + unframed + [pad, edge]

print(
    '\n'.join(add_border(['abc', 'ded'])),
    end='\n\n')

print('\n'.join(add_border([
        "Let's try a",
        "big sentence",
        "instead of these",
        "dumb letters."])),
      end='\n\n')

print('\n'.join(
    add_border(['Eat', 'mor', 'chikin'])),
    end='\n\n')


import os
import sys
import argparse
import logging
from pyspark.sql.functions import col, when
from pyspark.sql import SparkSession
from table_mappings import get_schema_struct, get_column_name_diff, get_primary_key

def _get_config(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', dest='host', type=str)
    parser.add_argument('--port', dest='port', type=int)
    parser.add_argument('--username', dest='username', type=str)
    parser.add_argument('--password', dest='password', type=str)
    parser.add_argument('--database', dest='database', type=str)
    parser.add_argument('--table', dest='table', type=str)
    parser.add_argument('--source', dest='source', type=str)
    parser.add_argument('--delimiter', dest='delimiter', type=str, default=',')
    config, _ = parser.parse_known_args(argv)

    return config

def _get_source_path(source):
    """Build absolute source file path

    :param source Source path
    """
    current_path = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(current_path, source)

def _write_to_db(data, config):
    return  data.write \
                .format("jdbc") \
                .option("url", f"jdbc:postgresql://{config.host}:{config.port}/{config.database}") \
                .option("dbtable", config.table) \
                .option("user", config.username) \
                .option("password", config.password) \
                .option("driver", "org.postgresql.Driver") \
                .mode('append') \
                .save()

def _map_columns(data, schema):
    return data.select([col(c).alias(schema.get(c, c)) for c in data.columns])

def _read_csv(spark, source_path, schema, delimiter=','):
    return spark.read \
                .format("csv") \
                .option("header", True) \
                .option("delimiter", delimiter) \
                .schema(schema) \
                .load(source_path) \
                .na.replace('NA', None)

def main(argv, spark):
    try:
        config = _get_config(argv)
        source_path = _get_source_path(config.source)

        raw_data = _read_csv(spark, source_path, get_schema_struct(config.table))
        mapped_data = _map_columns(raw_data, get_column_name_diff(config.table))

        if config.table == "flights":
            mapped_data = mapped_data.drop("index")
            #mapped_data_with_delay = mapped_data.withColumn("is_delayed", when((col("dep_delay") > 15), True).otherwise(False))

        return _write_to_db(mapped_data.dropDuplicates(get_primary_key(config.table)), config)

    except Exception as e:
        return logging.exception(e)

if __name__ == '__main__':
    spark = SparkSession \
        .builder \
        .master("local") \
        .config("spark.jars", "C:\spark\spark-3.4.0-bin-hadoop3\jars\postgresql-42.6.0.jar") \
        .appName("load_dw") \
        .getOrCreate()

    logging.getLogger().setLevel(logging.INFO)
    main(sys.argv, spark)


from ProjectUtils import *
from gym import spaces
from gym.utils import seeding

import numpy as np
from math import floor
import gym

from keras.models import Sequential, Model
from keras.layers import Dense, Activation, Flatten, Input, Concatenate, LSTM
from keras.optimizers import Adam
from pandas.plotting import register_matplotlib_converters

register_matplotlib_converters()
from rl.agents import DDPGAgent
from rl.memory import SequentialMemory
from rl.random import OrnsteinUhlenbeckProcess

ENV_NAME = "TEST"

""" LOAD DATA """

# df = pd.read_pickle("data/crypto_dat_v4.pickle", )
df = pd.read_csv('crypto_dat_v6.csv')

# print(df)

# df['timestamp'] = df['dates']
# df = df.drop('dates',axis=1)
df.fillna(0)
df = df[np.isfinite(df['Weighted_Price'])]

neuron_multiplier = 10

# TODO reorg max and min length to be in a number of days and not a fraction
days_length_of_training_episode = 30
training_range_percent = 0.775
print(len(df))
training_stop_index = int(len(df) * training_range_percent)
max_training_session_length = len(df) - training_stop_index

# df = df[500:]

class Stocks(gym.Env):

    def __init__(self, pandas_df):
        plt.style.use('dark_background')

        """ PROCESS DATA """
        other_data_keys = ['corrected_data','High','Low','Volume_(BTC)','Volume_(Currency)','Open','Close']
        self.price_keys = ['Weighted_Price', ]
        self.internal_tracking_keys = ['Working Balance', 'Bitcoin Balance', 'Bank Balance', 'Total Worth', 'Transferable Money', 'previous_Weighted_Price', 'previous_prediction_results', 'Baseline']

        # if isinstance(type(df['Timestamp'][0]),type('')):
        #     self.times_df = pd.to_datetime(df['Timestamp'])
        # else:
        self.times_df = pd.to_datetime(df['Timestamp'])

        input_columns = other_data_keys + self.price_keys + self.internal_tracking_keys

        self.df_inputs = pandas_df.reindex(columns=input_columns)

        # self.df_inputs = self.df_inputs.drop('Timestamp', axis=1)

        # print(self.df_inputs.columns)
        self.df_inputs.loc[-1] = np.zeros((len(self.df_inputs.columns.to_list())))
        # self.inputs, self.norms = norm_ops(self.df_inputs)
        self.inputs, self.norms = take_zscore(self.df_inputs)
        self.inputs = self.inputs.to_numpy()  # convert to n
        # umpy array
        # print(self.inputs)
        # print(self.inputs)

        """ DEFINE ACTION AND OBSERVATION SPACE """
        # Set number of actions (buy/sell for coins as well as hold/transfer)
        self.range = 1
        self.num_actions = len(self.price_keys) * 2 + 2
        self.action_space = spaces.Box(low=np.zeros((self.num_actions)), high=np.ones((self.num_actions)))  # , shape=(self.num_actions,), dtype=np.float32)
        self.observation_space = spaces.Box(low=0, high=self.range, shape=[self.df_inputs.shape[1]], dtype=np.float32)
        self.action_space.n = self.num_actions

        """ DEFINE A DICTIONARY TO FIND CORRECT INPUTS BY NAME LATER ON """
        self.key_df_column_dict = {}
        for i in range(len(input_columns)):
            self.key_df_column_dict.update({input_columns[i]: self.df_inputs.columns.get_loc(input_columns[i])})

        """ DEFINE UNCHANGING ENV VARS """
        self.trans_fee = 1
        self.transaction_fee = self.norm(self.trans_fee)
        self.total_step_counter = 0

        """ SET VARS FOR RENDER FUNCTION """
        self.render_keys = self.price_keys + ['Working Balance', 'Bitcoin Balance', 'Bank Balance', 'Total Worth', 'Transferable Money', 'Baseline']  # self.internal_tracking_keys #[:len(self.internal_tracking_keys)-2]

        # find what key corresponds to which index for render fnc

        self.render_df = self.create_empty_render_df()
        self.colours = ['xkcd:ultramarine',
                        'xkcd:vivid purple',
                        'xkcd:gold',
                        'xkcd:irish green',
                        'xkcd:dark pink',
                        'xkcd:lighter green',
                        'xkcd:sea blue']
        self.done = False

        """ START ENV """
        self.seed()

        # print(input_columns)
        # print(self.inputs)

    """ CORE FUNCTIONS """

    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    def step(self, action):
        """ MOST IMPORTANT FUNCTION FOR TIMING"""
        # print(action)
        """ PARSE ACTION """
        assert type(action) == np.ndarray
        action = list(action)
        # for i in range(len(action)):
        #     if action[i] < 0:
        #         action[i]= 0
        #     else:
        #         pass
        # action = [abs(x) for x in action]
        # print(action)
        self.chosen_action = int(np.argmax(action))
        chosen_action_amt = action[self.chosen_action]

        """ REWARD """
        # reward has 4 terms:
        # 1. the total worth compared to initial worth
        # 2. the amount in the bank compared to initial worth
        # 3. a punishment factor for not picking a new action
        # 4. a constant to shift the rewards into a range near 0 at the start
        # 5. a small bonus for making a new choice
        self.reward = 0# self.denorm(self.total_worth) / self.denorm(self.baseline_worth) + (self.denorm(self.bank) / self.denorm(self.random_balance)) - 1  # - np.exp(self.same_actions_counter * 0.01) + 1
        # if self.same_actions_counter ==0:
        #     self.reward +=  np.exp(0.1)
        # print(' same action = ', self.same_actions_counter )

        """ UPDATE STATE """
        coin_values = []
        for i in range(len(self.coin_amounts)):
            coin_values.append(self.inputs[self.step_index][self.df_inputs.columns.get_loc(self.price_keys[i])])

        self.update_internal_data()
        self.process_action(self.chosen_action, chosen_action_amt, coin_values)
        self.time_rewards()
        self.step_index += 1

        """ DEFINE HOW ACTIONS WORK TO AGENT """
        coin_values = []
        for i in range(len(self.coin_amounts)): coin_values.append(self.inputs[self.step_index][self.df_inputs.columns.get_loc(self.price_keys[i])])

        for i in range(len(self.price_keys)):  self.coin_values[i] = self.norm(self.coin_amounts[i] * self.denorm(coin_values[i]))

        self.baseline_worth = self.norm(self.baseline_amount * self.denorm(coin_values[0]))

        """ ADD INTERNAL DATA TO NEXT INPUT """

        self.state = self.inputs[self.step_index]
        # self.assess_if_done()
        # self.save()

        return self.state, self.reward, self.done, {}

    def reset(self):
        self.new_start()
        return self.state

    def render(self, mode='live', title=None, **kwargs):

        plt.ion()
        plt.clf()
        self.lists_to_plot = []
        for i in range(len(self.render_keys)):
            self.lists_to_plot.append([])
            self.lists_to_plot[i] = self.render_df[self.render_keys[i]].iloc[:self.step_index].to_list()
        self.ylists = self.lists_to_plot
        self.xlist = self.times_df.iloc[:self.step_index].to_list()
        for i in range(len(self.render_keys)):  # Graph stuff
            graph = plt.plot(self.xlist, self.ylists[i], label=self.render_keys[i], color=self.colours[i])[0]
            graph.set_ydata(self.ylists[i])
            graph.set_xdata(self.xlist)

        plt.xlim(self.xlist[self.render_index] - pd.Timedelta('30 days'), self.xlist[self.render_index])
        plt.legend(loc="best")
        plt.draw()
        plt.pause(.00001)

        self.render_index += 1

    """ CLASS UTILITY FUNCTIONS """

    def assess_if_done(self):

        if undo_single_zscore(self.bank, self.norms, self.key_df_column_dict[self.price_keys[0]]) < 0:
            self.done = True

        # if self.same_actions_counter >= 7:  # must make a new action every 1 week(s)
        #     self.done = True

    def new_start(self):
        self.step_index = 0# np.random.randint(0, training_stop_index)
        self.start_index = self.render_index = self.step_index
        self.random_balance = np.random.randint(1000, 20000)
        money_frac = np.random.random(1)[0]
        coin_frac = 1 - money_frac
        print('A new Start')

        self.initial_balance = self.random_balance * money_frac  # 10000  # self.df_inputs[self.price_keys[0]][self.step_index+3]  # - 2300

        self.balance = self.norm(self.initial_balance)
        self.start_balance = self.balance  # this is to track reward

        self.actions_taken_list = []
        self.same_actions_counter = 0
        self.action_index = 0
        self.bank = self.norm(0)
        frac_coin_splits = (self.random_balance * coin_frac) / len(self.price_keys)

        self.coin_amounts = np.zeros((len(self.price_keys)))
        self.coin_values = np.zeros((len(self.price_keys)))
        for i in range(len(self.coin_values)):
            self.coin_values[i] = self.norm(frac_coin_splits)
            self.coin_amounts[i] = frac_coin_splits / self.denorm(self.inputs[self.step_index][self.key_df_column_dict[self.price_keys[i]]])

        self.baseline_worth = self.norm(self.random_balance)
        self.baseline_amount = self.random_balance / self.denorm(self.inputs[self.step_index][self.key_df_column_dict[self.price_keys[0]]])

        self.done = False
        self.render_df = self.create_empty_render_df()

        self.transferable = self.norm(0)
        self.chosen_action = 0
        self.state = self.inputs[self.step_index]
        self.update_internal_data()

        # return self.state

    def update_internal_data(self):
        self.total_worth = self.norm(self.denorm(self.balance) + self.denorm(self.coin_values.sum()))

        self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[0]]] = self.balance
        self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[1]]] = self.coin_values[0]
        self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[2]]] = self.bank
        self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[3]]] = self.total_worth
        self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[4]]] = self.transferable

        try:
            self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[5]]] = self.inputs[self.step_index - 1][self.key_df_column_dict[self.internal_tracking_keys[5]]]
            self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[6]]] = self.inputs[self.step_index - 1][self.key_df_column_dict[self.internal_tracking_keys[5]]]
        except IndexError:
            self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[5]]] = self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[5]]]
            self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[6]]] = self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[5]]]

        self.inputs[self.step_index][self.key_df_column_dict[self.internal_tracking_keys[7]]] = self.baseline_worth

        """ ADD DATA TO RENDER FUNCTION """
        dat = []

        for i in range(len(self.render_keys)):
            dat.append(self.denorm(self.state[self.key_df_column_dict[self.render_keys[i]]]))
        self.render_df.iloc[self.step_index] = dat
        # print(self.render_df.iloc[self.step_index])

    def time_rewards(self):
        if self.action_index > 2:
            a = self.actions_taken_list[self.action_index - 1]
            b = self.actions_taken_list[self.action_index]
            if a == b:
                self.same_actions_counter += 1
            else:
                self.same_actions_counter = 0

        self.action_index += 1

    def denorm(self, x):
        return undo_single_zscore(x, self.norms, self.key_df_column_dict[self.price_keys[0]])

    def norm(self, x):
        return take_single_zscore(x, self.norms, self.key_df_column_dict[self.price_keys[0]])

    def process_action(self, chosen_act, chosen_amt, coin_vals):
        # print('transacion fees: ')
        # print(self.transaction_fee)
        # print(undo_single_zscore(self.transaction_fee,self.norms,self.df_inputs.columns.to_list().index(self.price_keys[0]), ))
        chosen_action_amt = chosen_amt
        chosen_action = chosen_act
        chosen_coin = 0
        has_coins = False
        has_money = False
        self.buy_action = False
        self.sell_action = False
        self.transfer_action = False
        self.money_transfered = False

        """ FIND TOTAL WORTH AND SEE IF ANY IS ALLOWED TO BE TRANSFERED """
        if (self.denorm(self.total_worth) / self.random_balance) > 1:
            self.transferable = self.norm(self.denorm(self.total_worth) - self.random_balance)
        if (self.denorm(self.total_worth) / self.random_balance) <= 1:
            self.transferable = self.norm(0)

        if chosen_action_amt < 0:  # actions can only be positive
            chosen_action_amt = -chosen_action_amt
        if chosen_action_amt > 1:
            chosen_action_amt = 1
        # print('chosen action = ',chosen_action)
        # print('chosen action amount = ',chosen_action_amt)

        if chosen_action <= len(self.coin_values) * 2:
            for i in range(len(self.coin_values) * 2):  # defining actions for buy/sell coins
                # which coin we're buying/selling
                if chosen_action % 2 == 0:  # if even, we're buying
                    chosen_coin = chosen_action / 2  # 0/2 = 0, 2/2 = 1
                    self.buy_action = True
                else:  # if odd, we're selling
                    chosen_coin = floor(chosen_action / 2)  # 1/2 = 0.5 floor-> 0, 3/2 = 1.5 floor -> 1
                    self.sell_action = True
                chosen_coin -= 1
            chosen_coin = int(chosen_coin)

        """ DENORMALIZED VALUES """
        coin_market_worth_denorm = self.denorm(coin_vals[chosen_coin])
        balance_denorm = self.denorm(self.balance)
        coin_val_denorm = self.coin_amounts[chosen_coin] * coin_market_worth_denorm
        # print(' num coins = ', self.coin_amounts[chosen_coin] )
        # print(' coin val = ', coin_market_worth_denorm)

        buy_amt = (balance_denorm * chosen_action_amt) - self.trans_fee
        sell_amt = (coin_val_denorm * chosen_action_amt)

        # print('coin_val_denorm ', coin_val_denorm)
        if coin_val_denorm > 0 and sell_amt <= (coin_val_denorm + self.trans_fee):
            has_coins = True
        if balance_denorm > self.trans_fee and buy_amt <= (balance_denorm + self.trans_fee) and buy_amt > 0:
            has_money = True
        if chosen_action == len(self.coin_values) * 2 + 1 and self.transferable > 0:
            self.transfer_action = True

        if self.buy_action and has_money:  # Buy (0 mod 2 = 0, 1 mod 2 = 1)
            self.balance = self.norm(balance_denorm - buy_amt)
            self.coin_amounts[chosen_coin] += buy_amt / coin_market_worth_denorm
            # print(' BUYING ' + str(buy_amt))

        elif self.sell_action and has_coins:  # Sell
            self.balance = self.norm(self.denorm(self.balance) + sell_amt)
            # print(sell_amt, coin_vals[chosen_coin])
            self.coin_amounts[chosen_coin] = self.coin_amounts[chosen_coin] - (sell_amt / coin_market_worth_denorm)
            # print(' SELLING ' + str(sell_amt))

        elif self.transfer_action and has_money:  # transfer to bank account
            # print(' TRANSFER ' + str(buy_amt))

            if buy_amt > self.denorm(self.transferable):
                buy_amt = self.denorm(self.transferable)
            self.bank = self.norm(self.bank + (buy_amt + self.trans_fee))  # add money to bank account
            self.balance = self.norm(balance_denorm - buy_amt)
            self.transferable = self.norm(self.denorm(self.transferable) - buy_amt)
            self.money_transfered = True

        else:  # hold
            # print(' HOLDING')
            pass

        self.actions_taken_list.append(chosen_action)

    def create_empty_render_df(self):
        render_df = pd.DataFrame({'dummy_dat': np.zeros((len(self.inputs)))})
        empty_df = render_df.reindex(columns=self.render_keys)

        # for i in range(len(self.render_keys)):
        #
        #     empty_df = render_df.reindex(columns=self.render_keys)#pd.DataFrame({self.render_keys[i]: np.empty((len(self.df_inputs))).tolist()})
        #
        #     if render_df.shape[0] == 0:  # if first time adding a column, use append
        #         render_df = render_df.append(empty_df)
        #
        #     else:  # if not, use concat for correct behaviour
        #         render_df = pd.concat((render_df, empty_df), sort=False, axis=1)

        return empty_df

    def save(self, save_on_step=100):

        self.total_step_counter += 1

        # print('totl steps = ',self.total_step_counter)
        if self.total_step_counter % save_on_step == 0 and self.total_step_counter > 0:
            agent.save_weights('{}_weights.h5'.format(ENV_NAME), overwrite=True)

# def make_model_and_stuff():
# Get the environment and extract the number of actions.
env = Stocks(df)
np.random.seed(123)
env.seed(123)
assert len(env.action_space.shape) == 1
nb_actions = env.action_space.shape[0]

actor_neurons = 4 * neuron_multiplier
actor_activation = 'relu'
# Next, we build a very simple model.
actor = Sequential()
actor.add(Flatten(input_shape=(1,) + env.observation_space.shape))
actor.add(Dense(actor_neurons))
actor.add(Activation(actor_activation))
actor.add(Dense(actor_neurons))
actor.add(Activation(actor_activation))
actor.add(Dense(nb_actions))
actor.add(Activation('linear'))
print(actor.summary())

critic_neurons = 8 * neuron_multiplier
critic_activation = 'relu'

action_input = Input(shape=(nb_actions,), name='action_input')
observation_input = Input(shape=(1,) + env.observation_space.shape, name='observation_input')
flattened_observation = Flatten()(observation_input)
x = Concatenate()([action_input, flattened_observation])
x = Dense(critic_neurons)(x)
x = Activation(critic_activation)(x)
x = Dense(critic_neurons)(x)
x = Activation(critic_activation)(x)
x = Dense(1)(x)
x = Activation('linear')(x)
critic = Model(inputs=[action_input, observation_input], outputs=x)
print(critic.summary())

# Finally, we configure and compile our agent. You can use every built-in Keras optimizer and
# even the metrics!
memory = SequentialMemory(limit=500000, window_length=1)
random_process = OrnsteinUhlenbeckProcess(size=nb_actions, theta=.15, mu=0.1, sigma=.3)

agent = DDPGAgent(nb_actions=nb_actions, actor=actor, critic=critic, critic_action_input=action_input,
                  memory=memory,  # nb_steps_warmup_critic=10, nb_steps_warmup_actor=10,
                  random_process=random_process, gamma=.99, target_model_update=1)

agent.compile(Adam(lr=0.001, clipnorm=1.), metrics=['mae'])

# Okay, now it's time to learn something! We visualize the training here for show, but this
# slows down training quite a lot. You can always safely abort the training prematurely using
# Ctrl + C.
# agent.load_weights('models/ddpg/{}_weights.h5f'.format(ENV_NAME))
print('Fitting agent.')
agent.fit(env, nb_steps=5000, visualize=True, verbose=1, nb_max_episode_steps=max_training_session_length, )
# agent.fit(env, nb_steps=500000, visualize=False, verbose=1, nb_max_episode_steps=max_training_session_length, )
# agent.fit(env, nb_steps=10000, visualize=True, verbose=1, nb_max_episode_steps=max_training_session_length, )

# After training is done, we save the final weights.
agent.save_weights('models/ddpg/{}_weights.h5f'.format(ENV_NAME), overwrite=True)

# out_df = pd.DataFrame({'agent':[agent]})
# save = open('ddpg_{}_agent.pickle'.format(ENV_NAME), 'wb')
# pickle.dump(out_df, save)
# save.close()
# pickle.dump(agent,'ddpg_{}_agent'.format(ENV_NAME))
# pickle.dump(agent,'ddpg_{}_agent'.format(ENV_NAME))

# Finally, evaluate our algorithm for 5 episodes.
# agent.test(env, nb_episodes=5, visualize=True, nb_max_episode_steps=200)


from setuptools import setup
from pybind11.setup_helpers import Pybind11Extension, build_ext

ext_modules = [
    Pybind11Extension(
        "harmonica/bindings",
        [
             'harmonica/orbit/kepler.cpp',
             'harmonica/orbit/trajectories.cpp',
             'harmonica/orbit/gradients.cpp',
             'harmonica/light_curve/fluxes.cpp',
             'harmonica/light_curve/gradients.cpp',
             'harmonica/bindings.cpp'
         ],
        include_dirs=["vendor/eigen", "vendor/pybind11"],
        language="c++",
        extra_compile_args=["-O2", "-ffast-math"]
    ),
]

setup(
    name="planet-harmonica",
    version="0.1.0",
    author="David Grant",
    author_email="david.grant@bristol.ac.uk",
    url="https://github.com/DavoGrant/harmonica",
    license="MIT",
    packages=["harmonica", "harmonica.jax"],
    description="Light curves for exoplanet transmission mapping.",
    long_description="Light curves for exoplanet transmission mapping.",
    python_requires=">=3.6",
    install_requires=["numpy", "jax", "jaxlib"],
    cmdclass={"build_ext": build_ext},
    ext_modules=ext_modules,
    classifiers=[
        "Intended Audience :: Science/Research",
        "License :: OSI Approved :: MIT License",
        "Operating System :: OS Independent",
        "Programming Language :: Python",
        "Programming Language :: Python :: 3",
        "Topic :: Scientific/Engineering :: Astronomy",
        "Topic :: Software Development :: Libraries :: Python Modules"
    ],
)


# 문제 출처 : https://www.acmicpc.net/problem/1021

from sys import stdin
from collections import deque

n, p = map(int, stdin.readline().split())
arr = list(map(int, stdin.readline().split()))
deque = deque([i for i in range(1, n+1)])

count = 0

for num in arr:
    if deque[0] == num:
        deque.popleft()
        continue

    left_move = deque.index(num)
    right_move = len(deque) - left_move

    if left_move <= right_move:
        deque.rotate(-left_move)
        deque.popleft()
        count += left_move
    else:
        deque.rotate(right_move)
        deque.popleft()
        count += right_move

print(count)


from configparser import ConfigParser
from multiprocessing import Pool
import os
import sys

if len(sys.argv) != 2:
    print('Arguments: config')
    sys.exit(-1)

cp = ConfigParser()
with open(sys.argv[1]) as fh:
    cp.read_file(fh)
cp = cp['config']
nb_classes = int(cp['nb_classes'])
dataset = cp['dataset']
random_seed = int(cp["random_seed"])
first_batch_size = int(cp["first_batch_size"])
il_states = int(cp["il_states"])
feat_root = cp["feat_root"]
incr_batch_size = (nb_classes-first_batch_size)//il_states

for state_id in range(il_states+1):
    print("Preparing state",state_id, "of", il_states)
    root_path = os.path.join(feat_root,"fetril",dataset,"seed"+str(random_seed),"b"+str(first_batch_size),"t"+str(il_states),"train","batch"+str(state_id))

    nb_classes = first_batch_size + (state_id) * incr_batch_size
    def decompose_class(n):
        file_path = os.path.join(root_path, str(n))
        if os.path.exists(file_path):
            try:
                os.makedirs(file_path+'_decomposed', exist_ok=True)
                compteur = 0
                with open(file_path, 'r') as f:
                    for line in f:
                        with open(os.path.join(file_path+'_decomposed', str(compteur)), 'w') as f2:
                            f2.write(line)
                        compteur += 1
            except:
                pass

    with Pool() as p:
        p.map(decompose_class, range(nb_classes))


import os
import cv2
import numpy as np
from keras.models import load_model
from keras.utils import to_categorical

# Define the path to the FER-2013 test dataset directory
test_data_path = "/content/Dissertation-Project/dataset/test"

# Define the list of emotions and their corresponding labels
emotion_labels = {
    'angry': 0,
    'disgust': 1,
    'fear': 2,
    'happy': 3,
    'neutral': 4,
    'sad': 5,
    'surprise': 6
}

# Load the test images and labels
test_images = []
test_labels = []

for emotion_folder in os.listdir(test_data_path):
    emotion = emotion_folder
    emotion_folder_path = os.path.join(test_data_path, emotion_folder)

    # Skip non-directory files
    if not os.path.isdir(emotion_folder_path):
        continue

    for image_filename in os.listdir(emotion_folder_path):
        # Skip non-image files
        if not image_filename.endswith(('.jpg', '.jpeg', '.png')):
            continue

        image_path = os.path.join(emotion_folder_path, image_filename)
        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        image = cv2.resize(image, (48, 48))  # Resize images to the desired shape
        test_images.append(image)
        test_labels.append(emotion_labels[emotion])

# Convert test images and labels to numpy arrays
test_images = np.array(test_images)
test_labels = np.array(test_labels)

# Normalize pixel values to range between 0 and 1
test_images = test_images / 255.0

# Perform one-hot encoding on the test labels
test_labels = to_categorical(test_labels)

# Load the trained model
model_path = "/content/Dissertation-Project/my_model_with_early_stopping.keras"
model = load_model(model_path)

# Evaluate the model on the test set
loss, accuracy = model.evaluate(test_images, test_labels)

print(f"Test Loss: {loss:.4f}")
print(f"Test Accuracy: {accuracy:.4f}")


import requests
from bs4 import BeautifulSoup
from time import sleep

movie_links_id = []

for z in range(10):
    print(z)

    url_1 = f"https://www.kinoafisha.info/rating/movies/?page={z}"
    sleep(5)
    r_1 = requests.get(url_1, timeout=5)
    soup_1 = BeautifulSoup(r_1.text, 'lxml')

    films = soup_1.find_all('div', class_='movieItem')
    for film in films:
        link = film.find('a', 'movieItem_title').get('href')
        link = ''.join([i for i in link if i.isdigit()])
        movie_links_id.append(link)


# basics
import argparse
import os
import pickle
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

# sklearn imports
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.neighbors import NearestNeighbors
from sklearn.preprocessing import normalize

# our code
import linear_model
import utils

url_amazon = "https://www.amazon.com/dp/%s"

def load_dataset(filename):
    with open(os.path.join('..','data',filename), 'rb') as f:
        return pickle.load(f)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-q','--question', required=True)
    io_args = parser.parse_args()
    question = io_args.question

    if question == "1":

        filename = "ratings_Patio_Lawn_and_Garden.csv"
        with open(os.path.join("..", "data", filename), "rb") as f:
            ratings = pd.read_csv(f,names=("user","item","rating","timestamp"))

        print("Number of ratings:", len(ratings))
        print("The average rating:", np.mean(ratings["rating"]))

        n = len(set(ratings["user"]))
        d = len(set(ratings["item"]))
        print("Number of users:", n)
        print("Number of items:", d)
        print("Fraction nonzero:", len(ratings)/(n*d))

        X, user_mapper, item_mapper, user_inverse_mapper, item_inverse_mapper, user_ind, item_ind = utils.create_user_item_matrix(ratings)
        print(type(X))
        print("Dimensions of X:", X.shape)

    elif question == "1.1":
        filename = "ratings_Patio_Lawn_and_Garden.csv"
        with open(os.path.join("..", "data", filename), "rb") as f:
            ratings = pd.read_csv(f,names=("user","item","rating","timestamp"))
        X, user_mapper, item_mapper, user_inverse_mapper, item_inverse_mapper, user_ind, item_ind = utils.create_user_item_matrix(ratings)
        X_binary = X != 0

        # YOUR CODE HERE FOR Q1.1.1
        items_groups = ratings.groupby(['item']).sum()
        items_groups['rating'].idxmax()
        items_groups.sort_values(by='rating',ascending=False)[0:5]

        # YOUR CODE HERE FOR Q1.1.2
        user_groups = ratings.groupby(['user']).size().reset_index(name='count')
        user_groups.sort_values(by='count',ascending=False)[0:5]

        # YOUR CODE HERE FOR Q1.1.3
        plt.hist(user_groups['count'], bins=45)
        plt.yscale('log', nonposy='clip')
        plt.xlabel('Number of ratings per user')
        plt.ylabel('log(Frequency)')
        plt.savefig('user_hist.png')

        items_groups = ratings.groupby(['item']).sum()
        plt.hist(items_groups['rating'],bins=40)
        plt.yscale('log', nonposy='clip')
        plt.xlabel('Number of ratings per item')
        plt.ylabel('log(Frequency)')
        plt.savefig('item_hist.png')

        plt.hist(ratings['rating'], bins=5)
        plt.xlabel('Rating out of 5 stars')
        plt.ylabel('rating_hist.png')

    elif question == "1.2":
        filename = "ratings_Patio_Lawn_and_Garden.csv"
        with open(os.path.join("..", "data", filename), "rb") as f:
            ratings = pd.read_csv(f,names=("user","item","rating","timestamp"))
        X, user_mapper, item_mapper, user_inverse_mapper, item_inverse_mapper, user_ind, item_ind = utils.create_user_item_matrix(ratings)
        X_binary = X != 0

        grill_brush = "B00CFM0P7Y"
        grill_brush_ind = item_mapper[grill_brush]
        grill_brush_vec = X[:,grill_brush_ind]

        print(url_amazon % grill_brush)

        # YOUR CODE HERE FOR Q1.2

        def nearest_ids(nbrs_list):
            nearest_id = np.empty(len(nbrs_list), dtype='S15')
            for key,value in item_inverse_mapper.items():
                    for i in range(0,len(nbrs_list)):
                        if key == nbrs_list[i]:
                            nearest_id[i] = value
            return nearest_id

        #Transpose so KNN iterates over columns not rows
        X_train = np.transpose(X)
        print(X_train.shape);print(X.shape)

        #Fit model
        model = NearestNeighbors(n_neighbors=6)
        model.fit(X_train)

        #Apply knn to get the index of the nearest items
        nbrs = model.kneighbors(np.transpose(grill_brush_vec), n_neighbors=6, return_distance=False)

        #get the item id's of the nearest items
        nbrs_idx = nbrs[0]
        nearest_id = nearest_ids(nbrs_idx)
        item_list = nearest_id[1:]
        print(item_list)

        #Normalize data and train model
        X_train_norm = normalize(X_train)
        model_norm = NearestNeighbors(n_neighbors=6)
        model_norm.fit(X_train_norm)

        #Knn to get nearest items from normned data
        nbrs_normed = model_norm.kneighbors(np.transpose(grill_brush_vec), n_neighbors=6, return_distance=False)

        #Get item id's of nearest items
        nbrs_idx_norm = nbrs_normed[0]
        nearest_norm_id = nearest_ids(nbrs_idx_norm)
        norm_item_list = nearest_norm_id[1:]
        print(norm_item_list)

        #Fit model based on cosine similarity
        model_cosine = NearestNeighbors(n_neighbors=6,metric='cosine')
        model_cosine.fit(X_train)

        #Knn to get nearest items
        nbrs_cosine = model_cosine.kneighbors(np.transpose(grill_brush_vec), n_neighbors=6, return_distance=False)

        #Get item id's of nearest items
        nbrs_idx_cosine = nbrs_cosine[0]
        nearest_cosine_id = nearest_ids(nbrs_idx_cosine)
        cosine_item_list = nearest_cosine_id[1:]
        print(cosine_item_list)

        #Are both lists actually the same
        print(cosine_item_list == norm_item_list)

        # YOUR CODE HERE FOR Q1.3

        list1 = [xi.decode() for xi in item_list]
        list2 = [xi.decode() for xi in cosine_item_list]
        list3 = [xi.decode() for xi in norm_item_list]
        print(list1);print(list2);print(list3)

        items_groups = ratings.groupby(['item']).size().reset_index(name='count')
        items_groups.sort_values(by='count',ascending=False)[0:5]

        count_L1 = [items_groups.loc[items_groups['item'] == xi] for xi in list1]
        count_L2 = [items_groups.loc[items_groups['item'] == xi] for xi in list2]
        print(count_L1);print(count_L2)

    elif question == "3":
        data = load_dataset("outliersData.pkl")
        X = data['X']
        y = data['y']

        # Fit least-squares estimator
        model = linear_model.LeastSquares()
        model.fit(X,y)
        print(model.w)

        utils.test_and_plot(model,X,y,title="Least Squares",filename="least_squares_outliers.pdf")

        #model2 =linear_model.WeightedLeastSquares()
        #z = np.ones(500)
        #model2.fit(X,y,z)
        #print(model2.w)

    elif question == "3.1":
        data = load_dataset("outliersData.pkl")
        X = data['X']
        y = data['y']

        z1 = np.ones(400)
        z2 = np.full((100),0.1)
        z = np.r_[z1,z2]

        # Fit weighted least-squares estimator
        model2 = linear_model.WeightedLeastSquares()
        model2.fit(X,y,z)
        print(model2.w)

        utils.test_and_plot(model2,X,y,title="Weighted Least Squares",filename="weighted_least_squares_outliers.png")

    elif question == "3.3":
        # loads the data in the form of dictionary
        data = load_dataset("outliersData.pkl")
        X = data['X']
        y = data['y']

        # Fit least-squares estimator
        model = linear_model.LinearModelGradient()
        model.fit(X,y)
        print(model.w)

        utils.test_and_plot(model,X,y,title="Robust (L1) Linear Regression",filename="least_squares_robust.pdf")

    elif question == "4":
        data = load_dataset("basisData.pkl")
        X = data['X']
        y = data['y']
        Xtest = data['Xtest']
        ytest = data['ytest']

        # Fit least-squares model
        model = linear_model.LeastSquares()
        model.fit(X,y)

        utils.test_and_plot(model,X,y,Xtest,ytest,title="Least Squares, no bias",filename="least_squares_no_bias.pdf")

    elif question == "4.1":
        data = load_dataset("basisData.pkl")
        X = data['X']
        y = data['y']
        Xtest = data['Xtest']
        ytest = data['ytest']

        #Linear regression with bias
        model = linear_model.LeastSquaresBias()
        model.fit(X,y)

        utils.test_and_plot(model,X,y,Xtest,ytest,title="Least Squares, bias",filename="least_squares_bias.pdf")

    elif question == "4.2":
        data = load_dataset("basisData.pkl")
        X = data['X']
        y = data['y']
        Xtest = data['Xtest']
        ytest = data['ytest']

        for p in range(11):
            print("p=%d" % p)
            model = linear_model.LeastSquaresPoly(p)
            model.fit(X,y)
            utils.test_and_plot(model,X,y,Xtest,ytest)

        #model = linear_model.LeastSquaresPoly(5)
        #model.fit(X,y)
        #utils.test_and_plot(model,X,y,Xtest,ytest,title="Poly fit",filename="poly_fit.pdf")

    else:
        print("Unknown question: %s" % question)

CPSC340 - Machine Learning/a3/code/main.py


# python -m pip install -r .\requirements.txt
# python .\merge.py .\annotatedPdfs\ final.pdf

import PyPDF2
import os
import sys

def merge_pdfs(pdf_dir, output_filename):
    pdf_writer = PyPDF2.PdfWriter()

    # Get all PDF files in the given directory
    files = [os.path.join(pdf_dir, file) for file in os.listdir(pdf_dir) if file.endswith('.pdf')]

    for file in files:
        with open(file, 'rb') as pdf_file:
            pdf_reader = PyPDF2.PdfReader(pdf_file)

            for page in range(len(pdf_reader.pages)):
                pdf_writer.add_page(pdf_reader.pages[page])

    with open(output_filename, 'wb') as merged_pdf:
        pdf_writer.write(merged_pdf)

if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python merge.py <PDF_DIRECTORY> <OUTPUT_FILENAME>")
        sys.exit(1)

    pdf_dir = sys.argv[1]
    output_filename = sys.argv[2]

    # Merge the PDFs into a single file
    merge_pdfs(pdf_dir, output_filename)


# Import necessary libraries
import openai
import sys
import json
import html
import re
import ssl
import os
import pprint
import nltk
import requests
import time

if not nltk.data.find('tokenizers/punkt'):
    nltk.download('punkt', quiet=True)

# Get the first command line argument
location = sys.argv[1]
sku = sys.argv[2]

# Initialize an empty dictionary for credentials
credentials = {}

# Define the file path to the credentials file
creds_file_path = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),  # Get the directory of the current file
    "../creds2.txt"  # Append the relative path to the credentials file
)
if os.path.exists('product.json'):
    os.remove('product.json')

# Define a class to represent a location
class Location:
    def __init__(self, website, user, city, phone, consumer_key, consumer_secret, api_key):
        self.website = website
        self.user = user
        self.city = city
        self.phone = phone
        self.consumer_key = consumer_key
        self.consumer_secret = consumer_secret
        self.api_key = api_key  # Here's the new attribute

# Initialize an empty list to store locations
locations = []

# Open the credentials file
with open(creds_file_path) as f:
    # Initialize variables for parsing the file
    website = None
    user = None
    city = None
    phone = None
    consumer_key = None
    consumer_secret = None
    openai.api_key = None
    # Loop over each line in the file
    for line in f:
        line = line.strip()  # Remove trailing and leading whitespace
        # If the line is a website (indicated by brackets), store it and reset other variables
        if line.startswith("[") and line.endswith("]"):
            if website and user and city and phone and consumer_key and consumer_secret and openai.api_key:
                locations.append(Location(website, user, city, phone, consumer_key, consumer_secret, openai.api_key))
            website = line[1:-1].lstrip()  # Remove the brackets and any leading whitespace
            user = None
            city = None
            phone = None
            consumer_key = None
            consumer_secret = None
            openai.api_key = None
        # If the line starts with a bracket but doesn't end with one, it's a multiline website; just store the first part
        elif line.startswith("["):
            website = line[1:]
        # If the line ends with a bracket but doesn't start with one, it's the end of a multiline website; append this part
        elif line.endswith("]"):
            website += line[:-1]
        # If the line contains " = ", it's a key-value pair; parse and store it
        elif website and " = " in line:
            key, value = line.split(" = ")
            if key == "user":
                user = value
            elif key == "city":
                city = value
            elif key == "phone":
                phone = value
            elif key.lower().endswith("_consumer_key"):
                consumer_key = value
            elif key.lower().endswith("_consumer_secret"):
                consumer_secret = value
            elif key == "openai.api_key":
                openai.api_key = value

    # Once we've parsed the entire file, check if there are any leftover variables and, if so, add another location
    if website and user and city and phone and consumer_key and consumer_secret and openai.api_key:
        locations.append(Location(website, user, city, phone, consumer_key, consumer_secret, openai.api_key))

# Print the locations
for location in locations:
    print("Using " + location.website + " as source product." + sku)
    base_url = "https://" + location.website + "/wp-json/wc/v3/products"
    city = location.city
    phone = location.phone
    consumer_key = location.website + "_consumer_key:" + location.consumer_key
    consumer_secret = location.website + "_consumer_secret:" + location.consumer_secret

    auth = (
     location.consumer_key,
     location.consumer_secret,
           )
    response = requests.get(f'{base_url}', auth=auth, params={'sku': sku})
    response.raise_for_status()

    if not response.json():
        print(f"No product found with SKU: {sku}")
        exit()

    product = response.json()[0]

    #time.sleep(1)
    # pprint.pprint(product)
    print("Source site: ", location.website)
    print("Product to clone: ", product['sku'], "Name: ", product['name'])
    print()
    break

time.sleep(1)
print("Turn AI loose...")
for location in locations[1:]:
    print("Using " + location.website + " as source product." + sku)
    base_url = "https://" + location.website + "/wp-json/wc/v3/products"
    consumer_key = location.website + "_consumer_key:" + location.consumer_key
    consumer_secret = location.website + "_consumer_secret:" + location.consumer_secret

    auth = (
     location.consumer_key,
     location.consumer_secret,
           )

    response = requests.get(f'{base_url}', auth=auth, params={'sku': sku})
    response.raise_for_status()

    if not response.json():
        print(f"No product found with SKU: {sku}")
        exit()

    product = response.json()[0]

    sku = product['sku']
    city = location.city
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages = [
            {
                "role": "system",
                "content": "You are a helpful budtender who knows all about the cannabis industry in Northern California. ",
                },
            {
                "role": "user",
                "content": f"I have a product with SKU '{sku}' named '{product['name']}' with a short description of '{product['short_description']}'."
                f"I need a new but similar name for this product that will both help with SEO and improve the product visibility in search engines."
                f"Use the city name '{city}' in the title. Examples of really good usage would be variations of things like Best in Burlingame or San Ramons best. "
                f"Don't stray too far from the core idea of the original name.  Add the city name to the product title somehow. "
                f"Limit the new product name to about 70 characters.  Do not use any punctuation or apostrophes or any single or double quotes. "
                f"Use proper capitalization. Optimize all for SEO.  Never use prices in the new titles."
                },
            ]
        )

    new_product_name = response['choices'][0]['message']['content'].strip()
    new_product_name = html.unescape(re.sub('<.*?>', '', new_product_name))
    print("Suggested new product name: ", new_product_name)
    product['name'] = new_product_name

    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages = [
            {
                "role": "system",
                "content": "You are a helpful budtender who knows all about the cannabis industry in and around the city of '{city}'. ",
                },
            {
                "role": "user",
                "content": f"I am a budtender looking to deliver cannabis products to my customers. "
                f"Create a list of 5 possible public meetup spots in the area of '{city}' that can be used as meetup spots if the customer doesnt want to meet at home or work. Prioritize discreet and easy access."
                f"The locations should be convenient, well known, and not near a school or police station. Present them as a bullet list with short descriptions as to why they are good meetup spots."
                },
            ]
        )

    meetup_spots = response['choices'][0]['message']['content'].strip()
    meetup_spots = html.unescape(re.sub('<.*?>', '', meetup_spots))
    product['description'] = product['description'] + " \n\n  Have your " + new_product_name +  " delivered to your home or work or choose one of these great meetup spots in " + city + "\n" +  meetup_spots
    #print("Suggested meetups spots for ",city,": ", meetup_spots)
    #print("New product description", product['description'])
    update_url = f'{base_url}/{product["id"]}'
    update_response = requests.put(update_url, json=product, auth=auth)
    update_response.raise_for_status()
    break