# -*- coding: utf-8 -*-
"""
delta_metrics
=============
Created on Mon May 21 09:00:01 2018

@author: Jon
"""
import numpy as np
import networkx as nx
import rivgraph.directionality as dy
import rivgraph.ln_utils as lnu

"""
This script contains algorithms that were ported from Matlab scripts provided
by Alejandro Tejedor to compute topologic and dynamic metrics on deltas. The
provided Matlab script required the bioinformatics toolbox; here we use
networkx to achieve the same result. Ported by Jon Schwenk.
The conversion was tested by computing metrics for the Wax Lake Delta
(provided by AT) and the Yenesei Delta (provided by JS)--perfect agreement
was found for all metrics, for both deltas, using both the original Matlab
scripts and the Python functions provided here.
JS has made some efficiency improvments to the code; otherwise most variable
names and code structure was matched to the original Matlab scripts.

Use at your own risk.
"""

def compute_delta_metrics(links, nodes):
    """Compute delta metrics."""
    # Delta metrics require a single apex node
    # This is not the ideal way to force a single inlet; adding the super-apex
    # is generally a much better approach. It has yet to be tested thoroughly,
    # though.
    links_m, nodes_m = ensure_single_inlet(links, nodes)

    # Ensure we have a directed, acyclic graph; also include widths as weights
    G = graphiphy(links_m, nodes_m, weight='wid_adj')

    if nx.is_directed_acyclic_graph(G) is not True:
        raise RuntimeError('Cannot proceed with metrics as graph is not acyclic.')

    # Compute the intermediate variables required to compute delta metrics
    deltavars = intermediate_vars(G)

    # Compute metrics
    metrics = dict()
    ner, pexc, ner_randomized = delta_nER(deltavars, N=200)
    metrics['nonlin_entropy_rate'] = ner
    metrics['nER_prob_exceedence'] = pexc
    metrics['nER_randomized'] = ner_randomized
    metrics['top_mutual_info'], metrics['top_conditional_entropy'] = top_entropy_based_topo(deltavars)
    metrics['top_link_sharing_idx'] = top_link_sharing_index(deltavars)
    metrics['n_alt_paths'] = top_number_alternative_paths(deltavars)
    metrics['resistance_distance'] = top_resistance_distance(deltavars)
    metrics['top_pairwise_dependence'] = top_s2s_topo_pairwise_dep(deltavars)
    metrics['flux_sharing_idx'] = dyn_flux_sharing_index(deltavars)
    metrics['leakage_idx'] = dyn_leakage_index(deltavars)
    metrics['dyn_pairwise_dependence'] = dyn_pairwise_dep(deltavars)
    metrics['dyn_mutual_info'], metrics['dyn_conditional_entropy'] = dyn_entropy_based_dyn(deltavars)

    return metrics

def ensure_single_inlet(links, nodes):
    """
    Ensure only a single apex node exists. This dumbly just prunes all inlet
    nodes+links except the widest one. Recommended to use the super_apex()
    approach instead if you want to preserve all inlets.

    All the delta metrics here require a single apex node, and that that node
    be connected to at least two downstream links. This function ensures these
    conditions are met; where there are multiple inlets, the widest is chosen.
    This function also ensures that the inlet node is attached to at least two
    links--this is important for computing un-biased delta metrics.
    The links and nodes dicts are copied so they remain unaltered; the altered
    copies are returned.

    """
    # Copy links and nodes so we preserve the originals
    links_edit = dict()
    links_edit.update(links)
    nodes_edit = dict()
    nodes_edit.update(nodes)

    # Find the widest inlet
    in_wids = []
    for i in nodes_edit['inlets']:
        linkid = nodes_edit['conn'][nodes_edit['id'].index(i)][0]
        linkidx = links_edit['id'].index(linkid)
        in_wids.append(links_edit['wid_adj'][linkidx])
    widest_inlet_idx = in_wids.index(max(in_wids))
    inlets_to_remove = nodes_edit['inlets'][:]

    # Remove inlet nodes and links until continuity is no longer broken
    badnodes = dy.check_continuity(links_edit, nodes_edit)
    if len(badnodes) > 0:
        raise RuntimeError('Provided (links, nodes) has source or sink at nodes: {}.'.format(badnodes))

    # Keep the widest inlet - delete all others (and remove their subnetworks)
    main_inlet = inlets_to_remove.pop(widest_inlet_idx)
    for i in inlets_to_remove:
        nodes_edit['inlets'].remove(i)
        badnodes = dy.check_continuity(links_edit, nodes_edit)
        while len(badnodes) > 0:
            badnode = badnodes.pop()
            # Remove the links connected to the bad node:
            # the hanging node will also be removed
            connlinks = nodes_edit['conn'][nodes_edit['id'].index(badnode)]
            for cl in connlinks:
                links_edit, nodes_edit = lnu.delete_link(links_edit,
                                                         nodes_edit, cl)

            badnodes = dy.check_continuity(links_edit, nodes_edit)

    # Ensure there are at least two links emanating from the inlet node
    conn = nodes_edit['conn'][nodes_edit['id'].index(main_inlet)]
    while len(conn) == 1:
        main_inlet_new = links_edit['conn'][links_edit['id'].index(conn[0])][:]
        main_inlet_new.remove(main_inlet)
        links_edit, nodes_edit = lnu.delete_link(links_edit, nodes_edit,
                                                 conn[0])

        # Update new inlet node
        nodes_edit['inlets'].remove(main_inlet)
        main_inlet = main_inlet_new[0]
        nodes_edit['inlets'] = nodes_edit['inlets'] + [main_inlet]
        conn = nodes_edit['conn'][nodes_edit['id'].index(main_inlet)]

    return links_edit, nodes_edit

def add_super_apex(links, nodes, imshape):
    """
    If multiple inlets are present, this creates a "super apex" that is
    directly upstream of all the inlet nodes. The synthetic links created
    have zero length and widths equal to the sum of the widths of the links
    connected to their respective inlet node.
    """

    # Get inlet nodes
    ins = nodes['inlets']

    if len(ins) <= 1:
        return links, nodes

    # Find the location of the super-apex by averaging the inlets' locations
    ins_idx = [nodes['idx'][nodes['id'].index(i)] for i in ins]
    rs, cs = np.unravel_index(ins_idx, imshape)
    apex_r, apex_c = np.mean(rs, dtype=int), np.mean(cs, dtype=int)
    apex_idx = np.ravel_multi_index((apex_r, apex_c), imshape)

    # Get the widths of the super-apex links -- these are just the summed
    # widths of all the links connected to each inlet node
    sa_widths = []
    for i in ins:
        lconn = nodes['conn'][nodes['id'].index(i)]
        sa_widths.append(sum([links['wid_adj'][links['id'].index(lid)] for lid in lconn]))

    # Add links from the super-apex to the inlet nodes
    # Widths are computed above
    # lengths are set to zero for these synthetic links
    for i, wid in zip(ins, sa_widths):
        in_idx = nodes['idx'][nodes['id'].index(i)]
        idcs = [apex_idx, in_idx]
        links, nodes = lnu.add_link(links, nodes, idcs)
        links['wid_adj'].append(wid)
        # we also append to the other attributes to keep fields the same length
        links['wid'].append(wid)
        links['wid_med'].append(wid)
        links['sinuosity'].append(wid)
        links['len'].append(0)
        links['len_adj'].append(0)

    # Add the super apex node field to the nodes dictionary
    # nodes = ln_utils.add_node(nodes, apex_idx, sa_lids)
    nodes['super_apex'] = nodes['id'][-1]

    return links, nodes

def delete_super_apex(links, nodes):
    """
    If you have a super apex, this function deletes it and connecting links.
    """

    # Get super apex node
    if 'super_apex' not in nodes:
        raise ValueError('no super apex detected.')

    # identify super apex
    super_apex = nodes['super_apex'][0]

    # identify connecting links
    super_links = nodes['conn'][nodes['id'].index(super_apex)]

    # delete links first
    for i in super_links:
        links, nodes = lnu.delete_link(links, nodes, i)

    # then delete super apex
    nodes = lnu.delete_node(nodes, super_apex, warn=True)

    return links, nodes

def graphiphy(links, nodes, weight=None, inletweights=None):
    """Converts RivGraph links and nodes into a NetworkX graph object.

    Converts the RivGraph links and nodes dictionaries into a NetworkX graph
    object.

    Parameters
    ----------
    links : dict
        RivGraph links and their attributes
    nodes : dict
        RivGraph nodes and their attributes
    weight : str, optional
        Link attribute to use to weight the NetworkX graph. If not provided or
        None, the graph will be unweighted (links of 1 and 0)
    inletweights : list, optional
        Optional manual weights for the inlets when using the super-apex
        functionality. Overrides the weight set by the inlet link attribute
        in favor of values from the provided list. List must be in the same
        order and have the same length as nodes['inlets'].

    Returns
    -------
    G : networkx.DiGraph
        Returns a NetworkX DiGraph object weighted by the link attribute
        specified in the optional parameter `weight`
    """
    if weight is not None and weight not in links.keys():
        raise RuntimeError('Provided weight key not in links dictionary.')

    if weight is None:
        weights = np.ones((len(links['conn']), 1))
    else:
        weights = np.array(links[weight])

    # Check weights
    if np.sum(weights <= 0) > 0:
        raise Warning('One or more of your weights is =< 0. This could cause problems later.')

    if inletweights is not None:
        if 'super_apex' not in nodes.keys():
            raise RuntimeError('Can only specify weights if super-apex has been added.')
        if len(inletweights) != len(nodes['inlets']):
            raise RuntimeError('graphiphy requires {} weights but {} were provided.'.format(len(nodes['inlets']), len(inletweights)))
        # Set weights of inlet links
        for inw, inl in zip(inletweights, nodes['inlets']):
            lconn = nodes['conn'][nodes['id'].index(inl)][:]
            lconn = [lc for lc in lconn if lc in nodes['conn'][nodes['id'].index(nodes['super_apex'])]]
            lidx = links['id'].index(lconn[0])
            weights[lidx] = inw

    G = nx.DiGraph()
    G.add_nodes_from(nodes['id'])
    for lc, wt, lid in zip(links['conn'], weights, links['id']):
        G.add_edge(lc[0], lc[1], weight=wt, linkid=lid)
    return G

def normalize_adj_matrix(G):
    """
    Normalize adjacency matrix.

    Normalizes a graph's adjacency matrix so the sum of weights of each row
    equals one. G is a networkx Graph with weights assigned.

    """
    # First, get adjacency matrix
    A = nx.to_numpy_array(G)
    # Normalize each node
    for r in range(A.shape[0]):
        rowsum = np.sum(A[r, :])
        if rowsum > 0:
            A[r, :] = A[r, :] / np.sum(A[r, :])

    return A

def intermediate_vars(G):
    """
    Compute interemediate variables and matrices.

    Computes the intermediate variables and matrices required to compute
    delta metrics. This function prevents the re-computation of many matrices
    required in the metric functions.

    """
    deltavars = dict()

    # The weighted adjacency matrix (A) of a Directed Acyclic Graph (DAG) has
    # entries a_{uv} that correspond to the fraction of the flux
    # present at node v that flows through the channel (vu). Flux partitioning
    # is done via channel widths.

    # Compute normalized weighted adjacency matrix
    A = normalize_adj_matrix(G)

    deltavars['A_w'] = A.T

    # Apex and outlet nodes
    deltavars['apex'], deltavars['outlets'] = find_inlet_outlet_nodes(deltavars['A_w'])

    """ Weighted Adj """
    deltavars['F_w'], deltavars['SubN_w'] = delta_subN_F(deltavars['A_w'])

    """ Weighted transitional"""
    deltavars['A_w_trans'] = np.matmul(deltavars['A_w'], np.linalg.pinv(np.diag(np.sum(deltavars['A_w'], axis=0))))
    deltavars['F_w_trans'], deltavars['SubN_w_trans'] = delta_subN_F(deltavars['A_w_trans'])

    """ Unweighted Adj"""
    deltavars['A_uw'] = np.array(deltavars['A_w'].copy(), dtype=bool)
    deltavars['F_uw'], deltavars['SubN_uw'] = delta_subN_F(deltavars['A_uw'])

    """ Unweighted transitional"""
    deltavars['A_uw_trans'] = np.matmul(deltavars['A_uw'], np.linalg.pinv(np.diag(np.sum(deltavars['A_uw'], axis=0))))
    deltavars['F_uw_trans'], deltavars['SubN_uw_trans'] = delta_subN_F(deltavars['A_uw_trans'])

    return deltavars

def find_inlet_outlet_nodes(A):
    """
    Find inlet and outlet nodes.

    Given an input adjacency matrix (A), returns the inlet and outlet nodes.
    The graph should contain a single apex
    (i.e. run ensure_single_inlet first).

    """
    apex = np.where(np.sum(A, axis=1) == 0)[0]
    if apex.size != 1:
        raise RuntimeError('The graph contains more than one apex.')
    outlets = np.where(np.sum(A, axis=0) == 0)[0]

    return apex, outlets

def compute_steady_state_link_fluxes(G, links, nodes, weight_name='flux_ss'):
    """Compute steady state fluxes through the network graph.

    Computes the steady state fluxes through links given a directed, weighted,
    NetworkX graph. The network should have only a single inlet (use either
    ensure_single_inlet() or add_super_apex() to do this). Additionally,
    this method will fail if the network has parallel edges. You should first
    run ln_utils artificial_nodes() function to break parallel edges, then
    re-compute link widths and lengths. Method after Tejedor et al 2015 [1]_.

    .. [1] Tejedor, Alejandro, et al. "Delta channel networks: 1. A
       graph‐theoretic approach for studying connectivity and steady state
       transport on deltaic surfaces."
       Water Resources Research 51.6 (2015): 3998-4018.

    Parameters
    ----------
    G : networkx.DiGraph
        NetworkX DiGraph object from graphiphy()
    links : dict
        RivGraph links dictionary
    nodes : dict
        RivGraph nodes dictionary
    weight_name : str, optional
        Name to give the new attribute in the links dictionary, is optional,
        if not provided will be 'flux_ss' (flux steady-state)

    Returns
    -------
    links : dict
        RivGraph links dictionary with new attribute
    """
    # Normalize the adjacency matrix
    An = normalize_adj_matrix(G)
    # Transposed adjacency required for computing F
    An_t = np.transpose(An)
    # Compute steady-state flux distribution
    fluxes, _ = delta_subN_F(An_t)

    # Fluxes are at-a-node and need to be translated to links
    fluxes = np.expand_dims(fluxes, 1)
    # Expand node-fluxes back to full adjacency matrix
    fw = fluxes * An
    # All nonzero elements in fw represent links where there is flux
    rows, cols = np.where(fw > 0)
    Gnodes = list(G.nodes)
    linkfluxes = np.zeros((len(links['id']), 1))  # Preallocate storage
    for (r, c) in zip(rows, cols):
        u = Gnodes[r]
        v = Gnodes[c]
        link_id = G.edges[u, v]['linkid']
        linkfluxes[links['id'].index(link_id)] = fw[r, c]

    # Store the fluxes in the links dict
    links[weight_name] = np.array(linkfluxes).flatten().tolist()

    return links

def delta_subN_F(A, epsilon=10**-10):
    """
    Compute steady state flux distribution.

    Computes the steady state flux distribution in the delta nodes when the
    system is fed with a constant unity influx from the Apex. Also defines the
    subnetworks apex-to-outlet.
    The SubN is an NxM matrix, where N is number of nodes and M is the number
    of outlets. For each mth outlet, its contributing subnetwork is given by
    the nonzero entries in SubN. The values in SubN are the degree of
    "belongingness" of each node to its subnetwork. If SubN(m,n) = 0, the m'th
    node does not belong to the n'th subnetwork; but if SubN(m,n) = 1, the m'th
    node belongs *only* to the n'th subnetwork. The values in SubN may be
    interpreted as the percentage of tracers that pass through node m that
    eventually make their way to the outlet of subnetwork n.

    """
    ApexID, OutletsID = find_inlet_outlet_nodes(A)

    """ Computing the steady-state flux, F """
    # To avoid boundary conditions and with the purpose of computing F, we
    # create a cycled version of the graph by connecting the outlet nodes
    # to the apex
    AC = A.copy()
    AC[ApexID, OutletsID] = 1

    # F is proportional to the eigenvector corresponding to the zero eigenvalue
    # of L=I-AC
    L = np.identity(AC.shape[0]) - np.matmul(AC,
                                             np.linalg.pinv(
                                                np.diag(np.sum(AC, axis=0))))
    d, v = np.linalg.eig(L)
    # Renormalize eigenvectors so that F at apex equals 1
    I = np.where(np.abs(d) < epsilon)[0]
    F = np.abs(v[:, I] / v[ApexID, I])

    """ Computing subnetworks """
    # R is null space of L(Gr)=Din(Gr-Ar(Gr)) - where Gr is the reverse graph,
    # Din the in-degree matrix, and Ar the adjacency matrix of Gr
    Ar = np.transpose(A)
    Din = np.diag(np.sum(Ar, axis=1))
    L = Din - Ar
    d, v = np.linalg.eig(L)
    # Renormalize eigenvectors to one
    for i in range(v.shape[1]):
        # set values below epsilon to 0
        v[:, i][v[:, i] < epsilon] = 0
        if np.max(v[:, i]) == 0:
            continue
        else:
            v[:, i] = v[:, i] / np.max(v[:, i])

    # Null space basis
    SubN = v[:, np.where(np.abs(d) < epsilon)[0]]
    I = np.where(SubN < epsilon)
    SubN[I[0], I[1]] = 0

    return np.squeeze(F), SubN

def nl_entropy_rate(A):
    """
    Compute nonlocal entropy rate.

    Computes the nonlocal entropy rate (nER) corresponding to the delta
    (inlcuding flux partition) represented by matrix A

    """
    # Compute steady-state flux and subnetwork structure
    F, SubN = delta_subN_F(A)
    F = F/np.sum(F)

    # Entropy per node
    Entropy = []
    for i in range(len(F)):
        I = np.where(SubN[i, :] > 0)[0]
        ent = -np.sum(SubN[i, I]*np.log2(SubN[i, I]))
        if len(I) > 1:
            ent = ent / np.log2(len(I))
        Entropy.append(ent)

    nER = np.sum(F*np.array(Entropy))

    return nER

def delta_nER(deltavars, N=500):
    """
    Compute nonlocal entropy rate.

    Compute the nonlocal entrop rate (nER) corresponding to the delta
    (including flux partition) represented by adjacency matrix A, and compares
    its value with the nER resulting from randomizing the flux partition.

    Returns
    -------
    pExc :
        the probability that the value of nER for a randomization of the fluxes
        on the topology dictated by A exceeds the actual value of nER. If the
        value of pExc is lower than 0.10, we considered that the actual partition
        of fluxes is an extreme value of nER
    nER_Delta :
        the nonlinear entropy rate for the provided adjacency matrix
    nER_randA :
        the nonlinear entropy rates for the N randomized deltas

    """
    A = deltavars['A_w_trans'].copy()
    nER_Delta = nl_entropy_rate(A)

    nER_randA = []
    for i in range(N):
        A_rand = A.copy()
        I = np.where(A_rand > 0)
        rand_weights = np.random.uniform(0, 1, (1, len(I[0])))
        A_rand[I] = rand_weights
        A_rand = np.matmul(A_rand, np.linalg.pinv(np.diag(np.sum(A_rand, axis=0))))
        nER_randA.append(nl_entropy_rate(A_rand))

    pExc = len(np.where(nER_randA > nER_Delta)[0]) / len(nER_randA)

    return nER_Delta, pExc, nER_randA

def top_entropy_based_topo(deltavars, epsilon=10**-10):
    """
    Compute topologic mutual information and conditional entropies.

    Computes the Topologic Mutual Information (TMI) and the Topologic
    Conditional Entropy for each subnetwork.

    """
    outlets = deltavars['outlets']

    # Fluxes at each node F and subnetworks subN
    F = deltavars['F_uw_trans'].copy()
    F = F/np.sum(F)
    SubN = deltavars['SubN_uw_trans'].copy()

    # Fluxes at links
    L_F = np.matmul(deltavars['A_uw_trans'], np.diag(F))

    TMI = np.empty((SubN.shape[1], 2))
    TCE = np.empty((SubN.shape[1], 2))
    for i in range(SubN.shape[1]):

        # Nodes that belong to subnetwork i
        nodes_in = np.where(SubN[:, i] > epsilon)[0]
        # Nodes that don't belong to subnetwork i
        nodes_out = np.where(SubN[:, i] < epsilon)[0]
        outlet_SubN = list(set(outlets).intersection(set(nodes_in)))[0]
        # Fluxes within subnetwork i - remove nodes_out
        subN_F = L_F.copy()
        subN_F[:, nodes_out] = 0
        subN_F[nodes_out, :] = 0

        # Compute fluxes leaving (Fn_out) and entering (Fn_in) each node in
        # the subnetwork, and total flux in the subnetwork (FS)
        Fn_out = np.sum(subN_F, axis=0)
        Fn_in = np.sum(subN_F, axis=1)
        FS = np.sum(subN_F)

        # Normalize all fluxes by FS
        subN_F = subN_F / FS
        Fn_out = Fn_out / FS
        Fn_in = Fn_in / FS

        # Compute TMI and TCE
        TMI_sum = 0
        TCE_sum = 0
        for ni in nodes_in:
            downN = np.where(subN_F[:, ni] > 0)[0]
            if len(downN) != 0:
                for d in downN:
                    TMI_sum = TMI_sum + subN_F[d, ni] * \
                              np.log2(subN_F[d, ni] / (Fn_in[d] * Fn_out[ni]))
                    TCE_sum = TCE_sum - subN_F[d, ni] * \
                              np.log2(subN_F[d, ni] * \
                              subN_F[d, ni] / Fn_in[d] / Fn_out[ni])
        TMI[i, 0] = outlet_SubN
        TMI[i, 1] = TMI_sum
        TCE[i, 0] = outlet_SubN
        TCE[i, 1] = TCE_sum

    return TMI, TCE

def top_link_sharing_index(deltavars, epsilon=10**-10):
    """
    Compute the link sharing index.

    Computes the Link Sharing Index (LSI) which quantifies the overlapping
    (in terms of links) of each subnetwork with other subnetworks in the
    delta.
    """
    outlets = deltavars['outlets']

    # Don't need weights
    A = deltavars['A_uw'].copy()

    # Set of links in the network (r, c)
    r, c = np.where(A==True)
    NL = r.shape[0]
    LinkBelong = np.zeros((NL, 1))

    # SubN indicates which nodes belong to each subnetwork
    SubN = deltavars['SubN_uw'].copy()
    NS = SubN.shape[1]
    SubN_Links = [[] for sl in range(NS)]

    # Evalueate LinkBelong and SubN_Links
    for i in range(NL):
        for k in range(NS):
            if SubN[r[i], k] > 0 and SubN[c[i], k] > 0:
                LinkBelong[i] = LinkBelong[i] + 1
                SubN_Links[k].append(i)

    # LSI is defined for each subnetwork as one minus the average
    # inverse LinkBelong
    LSI = np.empty((NS, 2))
    for k in range(NS):
        I = np.where(SubN[outlets, k] > epsilon)[0]
        LSI[k, 0] = outlets[I]
        LSI[k, 1] = 1 - np.nanmean(1 / LinkBelong[SubN_Links[k]])

    return LSI

def top_number_alternative_paths(deltavars, epsilon=10**-15):
    """
    Compute number of alternative paths.

    Computes the number of alternative paths (Nap) in the combinatorics sense
    from the Apex to each of the shoreline outlets.

    """
    apexid = deltavars['apex']
    outlets = deltavars['outlets']

    # Don't need weights
    A = deltavars['A_uw'].copy()

    # To compute Nap we need to find the null space of L==I*-A', where I* is
    # the Identity matrix with zeros for the diagonal entries that correspond
    # to the outlets.
    D = np.ones((A.shape[0], 1))
    D[outlets] = 0
    L = np.diag(np.squeeze(D)) - A.T
    d, v = np.linalg.eig(L)
    d = np.abs(d)
    null_space_v = np.where(np.logical_and(d < epsilon, d > -epsilon))[0]

    # Renormalize eigenvectors of the null space to have one at the outlet entry
    vN = np.abs(v[:, null_space_v])
    paths = np.empty((null_space_v.shape[0], 2))
    for i in range(null_space_v.shape[0]):
        I = np.where(vN[outlets, i] > epsilon)[0]
        vN[:, i] = vN[:, i] / vN[outlets[I], i]
        paths[i, 0] = outlets[I]
        paths[i, 1] = vN[apexid, i]

    return paths

def top_resistance_distance(deltavars, epsilon=10**-15):
    """
    Compute the topologic resistance distance.

    NOTE! TopoDist was not supplied with this function--can use networkX to
    compute shortest path but need to know what "shortest" means
    This function will not work until TopoDist is resolved.
    Computes the resistance distance (RD) from the Apex to each of the
    shoreline outlets. The value of RD between two nodes is the effective
    resistance between the two nodes when each link in the network is replaced
    by a 1 ohm resistor.
    """

    apexid = deltavars['apex']
    outlets = deltavars['outlets']

    # Don't need weights
    As = deltavars['A_uw'].copy()

    # Compute the RD within each subnetwork
    SubN = deltavars['SubN_w'].copy()

    RD = np.empty((SubN.shape[1], 2))
    for i in range(SubN.shape[1]):
        # Nodes that don't belong to subnetwork
        I = np.where(np.abs(SubN[:, i]) < epsilon)[0]
        # Zero columns and rows of nodes that are not present in subnetwork i
        As_i = As.copy()
        As_i[I, :] = 0
        As_i[:, I] = 0
        # Laplacian L and its pseudoinverse
        L = np.diag(np.sum(As_i, axis=0)) - As_i
        invL = np.linalg.pinv(L)

        # Compute RD
        I = np.where(SubN[outlets, i] > epsilon)[0]
        o = outlets[I]
        a = apexid
        RD[i, 0] = o

        # Distance between the apex and the ith outlet
        TopoDist = graphshortestpath(As_i, a[0], o[0])

        # RD is normalized by TopoDist to be able to compare networks of different size
        RD[i, 1] = (invL[a, a] + invL[o, o] - invL[a, o] - \
                   invL[o, a]) / TopoDist

    return RD

def graphshortestpath(A, start, finish):
    """
    Find the shortest path.

    Uses networkx functions to find the shortest path along a graph defined
    by A; path is simply defined as the number of links. Actual length not
    considered. Number of links in the shortest path is returned.

    """
    G = nx.from_numpy_array(A)
    sp = nx.shortest_path_length(G, start, finish)

    return sp

def top_s2s_topo_pairwise_dep(deltavars, epsilon=10**-10):
    """
    Compute subnetwork topologic pairwise dependence.

    This  function computes the Subnetwork to Subnetwork Topologic Pairwise
    Dependence (TPD) which quantifies the overlapping for all pairs of
    subnetworks in terms of links.

    """
    outlets = deltavars['outlets']

    # Don't need weights
    A = deltavars['A_uw'].copy()

    # Set of links
    r, c = np.where(A > 0)
    NL = len(r)

    # SubN indicates which nodes belong to each subnetwork
    SubN = deltavars['SubN_uw'].copy()

    NS = SubN.shape[1]
    SubN_Links = [[] for i in range(NS)]

    # Evaluate SubN_Links
    for i in range(NL):
        for k in range(NS):
            if SubN[r[i], k] > 0 and SubN[c[i], k] > 0:
                SubN_Links[k].append(i)

    # Compute TDP
    TDP = np.empty((len(outlets), len(outlets)))
    for i in range(NS):
        for k in range(NS):
            TDP[k, i] = len(set(SubN_Links[i]).intersection(set(SubN_Links[k]))) / len(SubN_Links[k])

    return TDP

def dyn_flux_sharing_index(deltavars, epsilon=10**-10):
    """
    Compute the flux sharing index.

    Computes the Flux Sharing Index (LSI) which quantifies the overlapping
    (in terms of flux) of each subnetwork with other subnetworks in the
    delta.

    """
    outlets = deltavars['outlets']

    # Set of links in the network (r, c)
    r, c = np.where(deltavars['A_w'] > 0)
    NL = r.shape[0]

    # SubN indicates which nodes belong to each subnetwork
    SubN = deltavars['SubN_w'].copy()

    NS = SubN.shape[1]
    SubN_Links = [[] for sl in range(NS)]

    # Evalueate SubN_Links
    for i in range(NL):
        for k in range(NS):
            if SubN[r[i], k] > 0 and SubN[c[i], k] > 0:
                SubN_Links[k].append(i)

    # FSI is defined for each subnetwork as one minus the average inverse SubN
    FSI = np.empty((NS, 2))
    for k in range(NS):
        I = np.where(SubN[outlets, k] > epsilon)[0]
        if len(I) != 0:
            FSI[k, 0] = outlets[I]
            # Downstream nodes of all the links in the subnetwork
            NodesD = r[SubN_Links[k]]
            FSI[k, 1] = 1 - np.nanmean(SubN[NodesD, k])
        else:
            FSI[k, 0] = np.nan
            FSI[k, 1] = np.nan

    return FSI

def dyn_leakage_index(deltavars, epsilon=10**-10):
    """
    Compute the leakage index.

    Computes the LI which accounts for the fraction of flux in subnetwork i
    leaked to other subnetworks.

    """
    apexid = deltavars['apex']
    outlets = deltavars['outlets']

    A = deltavars['A_w'].copy()

    # Check that the inlet node is at a bifurcation
    a = apexid
    I = np.where(A[:, a] > 0)[0]
    if len(I) < 2:
        print('Warning: the apex of the delta has only one node downstream. It is recommended that there be at least two downstream links from the apex to avoid biases.')

    # Fluxes at each node F and subnetworks subN
    F = deltavars['F_w'].copy()
    SubN = deltavars['SubN_w'].copy()

    # Link fluxes
    L_F = np.matmul(deltavars['A_w_trans'], np.diag(F))

    # Mathematically LI is computed for each subnetwork as the difference on
    # the fluxes at the nodes minus the links normalized by the total flux
    # in the links
    LI = np.empty((SubN.shape[1], 2))
    for i in range(SubN.shape[1]):
        # Nodes that belong to subnetwork i
        nodes_in = np.where(SubN[:, i] > epsilon)[0]
        # Nodes that do not belong to subnetwork i
        nodes_out = np.where(SubN[:, i] < epsilon)
        outlet_subN = set(outlets).intersection(set(nodes_in))
        if len(outlet_subN) != 0:
            outlet_subN = outlet_subN.pop()
            # Fluxes within subnetwork i -- remove nodes_out
            subN_F = L_F.copy()
            subN_F[:, nodes_out] = 0
            subN_F[nodes_out, :] = 0
            # Active links within subnetwork
            links_subN = np.where(subN_F > 0)
            # Sum of the fluxes in all the li nks in the subnetwork
            sum_links = np.sum(subN_F[links_subN[0], links_subN[1]])
            # Sum of the fluxes in all the nodes (except the outlet--since it
            # cannot leak out by definition)
            sum_nodes = np.sum(F[nodes_in]) - F[outlet_subN]

            LI[i,0] = outlet_subN
            LI[i,1] = (sum_nodes - sum_links) / sum_nodes
        else:
            LI[i,:] = np.nan

    return LI

def dyn_pairwise_dep(deltavars, epsilon=10**-10):
    """
    Compute subnetwork dynamic pairwise dependence.

    Computes the subnetwork to subnetwork dynamic pairwise dependence (DPD)
    which quantifies the overlapping for all pairs of subnetworks in terms of
    flux.

    """
    A = deltavars['A_w_trans'].copy()

    # Set of links in the network (r, c)
    r, c = np.where(A > 0)
    NL = r.shape[0]

    # SubN indicates which nodes belong to each subnetwork
    F = deltavars['F_w_trans'].copy()
    SubN = deltavars['SubN_w_trans'].copy()

    NS = SubN.shape[1]
    SubN_Links = [[] for sl in range(NS)]

    # Evalueate SubN_Links
    for i in range(NL):
        for k in range(NS):
            if SubN[r[i], k] > 0 and SubN[c[i], k] > 0:
                SubN_Links[k].append(i)

    # Link fluxes
    L_F = np.matmul(A, np.diag(F))

    # Compute DPD
    DPD = np.empty((NS, NS))
    for i in range(NS):
        for k in range(NS):
            link_intersect = list(set(SubN_Links[i]).intersection(set(SubN_Links[k])))
            links_in_s = SubN_Links[k]
            DPD[k, i] = np.sum(L_F[r[link_intersect],
                               c[link_intersect]]) / np.sum(L_F[r[links_in_s],
                                                            c[links_in_s]])

    return DPD

def dyn_entropy_based_dyn(deltavars, epsilon=10**-10):
    """
    Compute dynamic mutual information and dynamic conditional entropy.

    Computes the Dynamic Mutual Information (DMI) and the Dynamic
    Conditional Entropy for each subnetwork.

    """
    outlets = deltavars['outlets']
    A = deltavars['A_w_trans'].copy()

    # Fluxes at each node F and subnetworks subN
    F = deltavars['F_w_trans'].copy()
    SubN = deltavars['SubN_w_trans'].copy()
    F = F / np.sum(F)
    # Fluxes at links
    L_F = np.matmul(A, np.diag(F))

    DMI = np.empty((SubN.shape[1], 2))
    DCE = np.empty((SubN.shape[1], 2))
    for i in range(SubN.shape[1]):

        # Nodes that belong to subnetwork i
        nodes_in = np.where(SubN[:, i] > epsilon)[0]
        # Nodes that don't belong to subnetwork i
        nodes_out = np.where(SubN[:, i] < epsilon)[0]
        outlet_SubN = list(set(outlets).intersection(set(nodes_in)))[0]
        # Fluxes within subnetwork i - remove nodes_out
        subN_F = L_F.copy()
        subN_F[:, nodes_out] = 0
        subN_F[nodes_out, :] = 0

        # Compute fluxes leaving (Fn_out) and entering (Fn_in) each node in
        # the subnetwork, and total flux in the subnetwork (FS)
        Fn_out = np.sum(subN_F, axis=0)
        Fn_in = np.sum(subN_F, axis=1)
        FS = np.sum(subN_F)

        # Normalize all fluxes by FS
        subN_F = subN_F / FS
        Fn_out = Fn_out / FS
        Fn_in = Fn_in / FS

        # Compute TMI and TCE
        DMI_sum = 0
        DCE_sum = 0
        for ni in nodes_in:
            downN = np.where(subN_F[:, ni] > 0)[0]
            if len(downN) != 0:
                for d in downN:
                    DMI_sum = DMI_sum + subN_F[d, ni] * \
                              np.log2(subN_F[d, ni] / (Fn_in[d] * Fn_out[ni]))
                    DCE_sum = DCE_sum - subN_F[d, ni] * \
                              np.log2(subN_F[d, ni] * subN_F[d, ni] / Fn_in[d] / Fn_out[ni])
        DMI[i, 0] = outlet_SubN
        DMI[i, 1] = DMI_sum
        DCE[i, 0] = outlet_SubN
        DCE[i, 1] = DCE_sum

    return DMI, DCE

def dist_from_apex(nodes, imshape):
    """Calculate normalized distance from apex.

    Does this for nodes. Calculates a normalized distances from apex, ignores
    pixel resolution.

    Parameters
    ----------
    nodes : dict
        RivGraph dictionary of nodes
    imshape : tuple
        Tuple of the shape of the domain (e.g., Imask.shape)

    Returns
    -------
    norm_dist : list
        List of normalized straight line distances between each node and the
        inlet in the same order as the nodes come in the input nodes
        dictionary.
    """
    # id row/coord of the apex (or representative location)
    apex_id = nodes['inlets']
    if len(apex_id) < 1:
        raise ValueError('No inlets')
    elif len(apex_id) > 1:
        # average inlet locations to a single point
        ins_idx = [nodes['idx'][nodes['id'].index(i)] for i in apex_id]
        rs, cs = np.unravel_index(ins_idx, imshape)
        apex_xy = np.mean(rs, dtype=int), np.mean(cs, dtype=int)
    else:
        apex_idx = nodes['idx'][nodes['id'].index(apex_id)]
        apex_xy = np.unravel_index(apex_idx, imshape)

    # calculate distances to all nodes from apex location
    def calc_dist(apex_xy, node_xy):
        """Euclidean distance function."""
        return np.sqrt((apex_xy[0]-node_xy[0])**2 +
                       (apex_xy[1]-node_xy[1])**2)

    # get coordinates of all nodes in xy space
    node_xy = [np.unravel_index(i, imshape) for i in nodes['idx']]
    node_dists = [calc_dist(apex_xy, i) for i in node_xy]
    # normalize and return this normalized distance
    norm_dist = list(np.array(node_dists) / np.max(node_dists))

    return norm_dist

def calc_QR(links, nodes, wt='wid_adj', new_at='graphQR'):
    """Clunky solution (target for optimization) to get QR at bifurcations.

    QR is defined as the larger branch Q / smaller branch Q per
    Edmonds & Slingerland 2008 [2]_. This measure of flux partitioning at a
    bifurcation does not scale beyond bifurcations to trifurcations etc.
    The graph-based flux partitioning scheme also assumes flow is routed
    in a steady-state manner based on the width (or some other attribute)
    of the links in the network. Therefore the actual flux value doesn't
    matter, we can calculate QR as larger width / smaller width from the two
    branches as that will be the same as if we'd calculated the steady-state
    fluxes and taken their ratio.
    The function is written flexibly to allow one to assuming flux weighting
    by an attribute other than the link width if desired.

    .. warning::

      QR values calculated at nodes located at confluences, polyfurcations,
      or any other non-bifurcating location will be incorrect!

    .. [2] Edmonds, D. A., and R. L. Slingerland. "Stability of delta
       distributary networks and their bifurcations."
       Water Resources Research 44.9 (2008).

    Parameters
    ----------
    links : dict
        RivGraph links dictionary
    nodes : dict
        RivGraph nodes dictionary
    wt : str, optional
        String pointing to the link attribute to use when calculating ratios,
        optional, default is 'wid_adj' which is the adjusted link widths
    new_at : str, optional
        Name of the new attribute to add to the nodes dictionary, optional,
        default is 'graphQR' to indicate the graph calculated QR value
    Returns
    -------
    nodes : dict
        RivGraph dictionary with new_at attribute added
    """
    # check links for wt attribute
    if wt not in links.keys():
        raise ValueError('wt attribute not in the links dictionary')

    # set up list of zeros
    nodes[new_at] = np.zeros_like(nodes['id'], dtype=float)

    for i in range(len(nodes['id'])):
        # for bifurcations
        if len(nodes['conn'][i]) == 3:
            # get the 3 connected link ids
            link_ids = nodes['conn'][i]
            # get upstream node for each link, its "start" point
            link_starts = [links['conn'][links['id'].index(link_ids[0])][0],
                           links['conn'][links['id'].index(link_ids[1])][0],
                           links['conn'][links['id'].index(link_ids[2])][0]]
            # figure out which two links are the ones leaving this node
            # and get the width of each
            # (which controls the local flux partitioning anyway)
            if link_starts[0] == link_starts[1]:
                # check if 1st and 2nd match
                wid_1 = links[wt][links['id'].index(link_ids[0])]
                wid_2 = links[wt][links['id'].index(link_ids[1])]
            elif link_starts[0] == link_starts[2]:
                # check if 1st and 3rd match
                wid_1 = links[wt][links['id'].index(link_ids[0])]
                wid_2 = links[wt][links['id'].index(link_ids[2])]
            else:
                # then 2nd and 3rd must match
                wid_1 = links[wt][links['id'].index(link_ids[1])]
                wid_2 = links[wt][links['id'].index(link_ids[2])]

        # for inlets w/ only 2 connecting links
        elif nodes['id'][i] in nodes['inlets'] and len(nodes['conn'][i]) == 2:
            link_ids = nodes['conn'][i]
            wid_1 = links[wt][links['id'].index(link_ids[0])]
            wid_2 = links[wt][links['id'].index(link_ids[1])]

        # catch-all for other scenarios: QR will be -1
        else:
            wid_1 = -1
            wid_2 = 1

        # calculate and assign QR to the node of interest
        wid_big = np.max([wid_1, wid_2])
        wid_small = np.min([wid_1, wid_2])
        nodes[new_at][i] = wid_big / wid_small

    # coerce into list
    nodes[new_at] = list(nodes[new_at])

    # if junction angles are known, make non-bifurcation node QR values NaNs
    if 'jtype' in nodes.keys():
        for i in range(len(nodes['jtype'])):
            if nodes['jtype'][i] != 'b':
                nodes[new_at][i] = np.nan

    return nodes


"""
DB GAPS 4회 AlphaLab 전략

모멘텀 스코어 기반
"""

# It requires kiwooma package which is python wrapper module for kiwoom OPEN API made by me

from kiwooma import EasyAPI
from collections import OrderedDict
from scipy.optimize import minimize
import time, datetime
import pandas as pd
import numpy as np

stock_weights = np.zeros(6)
bond_weights = np.zeros(3)

def maximize_score(scores, cov):

    def obj_func(w):
        # total score 최대화 함수
        w = w.reshape(len(scores), 1)
        new_scores = scores.copy()
        # new_scores.iloc[0:6] = scores.iloc[0:6].mean()
        # new_scores.iloc[6:9] = scores.iloc[6:9].mean()
        total_score = -(w.T).dot(new_scores)
        return total_score

    def vol(w):
        return (w.T.dot(cov).dot(w) * 20)**(1/2)

    def allocate_stocks(w):
        global stock_weights
        new_scores = scores.iloc[0:6].copy()
        new_scores[new_scores<0] = 0.0
        weights = np.array(new_scores/new_scores.sum())*0.4
        stock_weights = weights
        return weights

    def allocate_bonds(w):
        global bond_weights
        new_scores = scores.iloc[6:9].copy()
        new_scores[new_scores<0] = 0.0
        weights = np.array(new_scores/new_scores.sum())*0.6
        bond_weights = weights
        return weights

    cons = ({
                'type': 'ineq',
                'fun': lambda w: 0.99 - w.sum()
            },
            {
                'type': 'ineq',
                'fun': lambda w: w - 0.0000001
            },
            # {
            #     'type': 'ineq',
            #     'fun': lambda w: 0.005 - vol(w)
            # },
            {   # Kodex200
                'type': 'ineq',
                'fun': lambda w: 0.4 - w[0]
            },
            {   # Kosdaq150
                'type': 'ineq',
                'fun': lambda w: 0.4 - w[1]
            },
            {   # Kodex200 & Kosdaq150
                'type': 'ineq',
                'fun': lambda w: (w[0] + w[1]) - 0.1
            },
            {   # Kodex200 & Kosdaq150
                'type': 'ineq',
                'fun': lambda w: 0.4 - (w[0] + w[1])
            },
            {   # 해외주식
                'type': 'ineq',
                'fun': lambda w: (w[2] + w[3] + w[4] + w[5]) - 0.1
            },
            {   # 해외주식
                'type': 'ineq',
                'fun': lambda w: 0.4 - (w[2] + w[3] + w[4] + w[5])
            },
            {   # 국채 10년
                'type': 'ineq',
                'fun': lambda w: 0.5 - w[6]
            },
            {   # 우량회사채
                'type': 'ineq',
                'fun': lambda w: 0.4 - w[7]
            },
            {   # 해외채권
                'type': 'ineq',
                'fun': lambda w: 0.4 - w[8]
            },
            {   # 해외채권
                'type': 'ineq',
                'fun': lambda w: w[8] - 0.05
            },
            {   # 채권
                'type': 'ineq',
                'fun': lambda w: 0.6 - (w[6] + w[7] + w[8])
            },
            {   # 채권
                'type': 'ineq',
                'fun': lambda w: (w[6] + w[7] + w[8]) - 0.2
            },
            {   # 채권
                'type': 'ineq',
                'fun': lambda w: allocate_bonds(w) - w[6:9]
            },
            {   # 금
                'type': 'ineq',
                'fun': lambda w: 0.15 - w[9]
            },
            {   # WTI
                'type': 'ineq',
                'fun': lambda w: 0.15 - w[10]
            },
            {   # 원자재
                'type': 'ineq',
                'fun': lambda w: 0.20 - (w[9] + w[10])
            },
            {   # 원자재
                'type': 'ineq',
                'fun': lambda w: (w[9] + w[10]) - 0.05
            },
            {   # KOSPI Inverse
                'type': 'ineq',
                'fun': lambda w: w[11]
            },
            {   # KOSPI Inverse
                'type': 'ineq',
                'fun': lambda w: (w[0]+w[1])-w[11]
            },
            {   # Dollar
                'type': 'ineq',
                'fun': lambda w: 0.2 - w[12]
            },
            {   # Dollar Inverse
                'type': 'ineq',
                'fun': lambda w: 0.2 - w[13]
            },
            {   # Dollar Inverse
                'type': 'ineq',
                'fun': lambda w: 0.2 - (w[12] + w[13])
            },
            {   # 단기자금
                'type': 'ineq',
                'fun': lambda w: 0.5 - w[14]
            },
            {   # 주식
                'type': 'ineq',
                'fun': lambda w: allocate_stocks(w) - w[0:6]
            },)

    #First guess: 동일가중
    w0 = np.ones(len(scores))/len(scores)

    res = minimize(obj_func, w0, constraints=cons, method = 'SLSQP') #Minimization
    res = (pd.Series(res['x'], index=scores.index)*100).round(4)

    global_stock = res.iloc[2:6].sum()
    bonds = res.iloc[6:9].sum()

    res.loc['S&P500'] = global_stock * stock_weights[2]/stock_weights[2:].sum()
    res.loc['EuroStoxx'] = global_stock * stock_weights[3]/stock_weights[2:].sum()
    res.loc['Nikkei'] = global_stock * stock_weights[4]/stock_weights[2:].sum()
    res.loc['CSI300'] = global_stock * stock_weights[5]/stock_weights[2:].sum()
    res.loc['국고채10년'] = bonds * bond_weights[0]/bond_weights.sum()
    res.loc['중기우량회사채'] = bonds * bond_weights[1]/bond_weights.sum()
    res.loc['하이일드'] = bonds * bond_weights[2]/bond_weights.sum()
    vol = vol(np.array(res))
    res['현금'] = 100-res.sum()
    return res, vol

def calc_scores(code, price):
    if code == '132030':
        returns_20 = price.iloc[-1]/price.iloc[-21]-1
        returns_40 = price.iloc[-1]/price.iloc[-41]-1
        score = (-returns_20 - ((1+returns_40)**0.5-1))/2
    elif code == '138230' or code == '139660':
        returns_40 = price.iloc[-1]/price.iloc[-41]-1
        returns_80 = price.iloc[-1]/price.iloc[-81]-1
        score = ((1+returns_40)**(1/2) + (1+returns_80)**(1/4)-2)/2

    elif code ==  '195930':
        score = 0
    else:
        returns_40 = price.iloc[-1]/price.iloc[-41]-1
        returns_80 = price.iloc[-1]/price.iloc[-81]-1
        returns_120 = price.iloc[-1]/price.iloc[-121]-1
        returns_160 = price.iloc[-1]/price.iloc[-161]-1
        score = ((1+returns_40)**(1/2) + (1+returns_80)**(1/4) + \
            (1+returns_120)**(1/6) + (1+returns_160)**(1/8)-4)/4

    return score

def get_data(api, code):
    df = api.get_daily_ohlcv(code, adj_close=1, repeat=1)
    price_df = df[['open', 'high', 'low', 'close']].mean(1)
    return price_df

def get_scores(price_dict, codes, iloc):
    scores = {}
    dfs = []

    for code in codes:
        price = price_dict[code].iloc[:iloc]

        returns = price.pct_change()
        dfs.append(returns)
        returns.name = code
        score = calc_scores(code, price)
        scores[code] = score
        time.sleep(0.2)
    df = pd.concat(dfs, axis=1)
    return scores, df

def main():

    stock_dict = OrderedDict({
        '069500': 'KODEX200', '232080': 'KOSDAQ150', '143850': 'S&P500',
        '195930': 'EuroStoxx', '238720': 'Nikkei', '192090': 'CSI300',
        '148070': '국고채10년', '136340': '중기우량회사채', '182490': '하이일드',
        '132030': '골드선물', '130680': '원유선물', '114800': 'KODEX인버스',
        '138230': '달러선물', '139660': '달러인버스', '130730': '단기자금'
    })
    api = EasyAPI()

    price_dict = {}
    for code in list(stock_dict.keys()):
        price_dict[code] = get_data(api, code)

    scores, df = get_scores(price_dict, list(stock_dict.keys()), -1)

    df.rename(columns = stock_dict, inplace=True)
    cov = df.iloc[-60:].cov()

    new_scores = {}

    for key, item in scores.items():
        new_scores[stock_dict[key]] = item

    ret, vol = maximize_score(pd.Series(new_scores), cov)

    # ret.loc['']
    date = datetime.datetime.today().date()
    with open('{} - scores.txt'.format(date), mode='a', encoding='utf-8') as f:
        f.write('Scores\n-----\n')
        f.write(str(new_scores))
        f.write('\nOur Scores\n-----\n')
        f.write(str(pd.Series(new_scores)*100))
        f.write('\n')
        f.write('Weights\n-----\n')
        f.write('\n')
        f.write(str(ret))
        f.write('\n')
        f.write('Volatiltiy\n-----\n')
        f.write(str(vol))
        f.write('\n')

if __name__ == '__main__':
    main()


import torch
from torch.nn.functional import leaky_relu, dropout
from torch.nn import Linear, ReLU, BatchNorm1d, ModuleList, L1Loss, LeakyReLU, Dropout, MSELoss,SELU
from torch_geometric.nn import Sequential, SAGEConv, Linear, to_hetero,HeteroConv,GATConv,GINEConv

class GNNEncoder(torch.nn.Module):
    def __init__(self,hidden_channels,out_channels,num_layers):
        super().__init__()

        self.hidden_channels = hidden_channels
        self.out_channels = out_channels
        self.num_layers = num_layers
        metadata = (['node1', 'node2', 'node3'],
                    [('node2', 'to', 'node3'),
                     ('node1', 'to', 'node2'),
                     ('node3', 'rev_to', 'node2'),
                     ('node2', 'rev_to', 'node1')])
        self.convs = torch.nn.ModuleList()
        layers = (self.hidden_channels, self.out_channels)
        for i in range(self.num_layers):
            conv = HeteroConv({
                edge_type: SAGEConv((-1, -1), layers[i])
                for edge_type in metadata[1]
            })
            self.convs.append(conv)

    def forward(self, x_dict, edge_index_dict):

        for conv in self.convs:
            x_dict = conv(x_dict, edge_index_dict)
            x_dict = {key: leaky_relu(x) for key, x in x_dict.items()}

        return x_dict

class EdgeDecoder(torch.nn.Module):
    def __init__(self, hidden_channels,out_channels):
        super().__init__()

        self.hidden_channels = hidden_channels
        self.out_channels = out_channels

        self.network = torch.nn.Sequential(
            Linear(-1, hidden_channels),
            SELU(),
            Dropout(p=0.2),
            Linear(hidden_channels, out_channels),
            SELU(),
            Dropout(p=0.2),
            Linear(out_channels, 32),
            SELU(),
            Dropout(p=0.2),
            Linear(32, 1)
        )

    def forward(self, z_dict, edge_index_dict,edge_attr_dict):
        ship, cust = edge_index_dict[('node1', 'to', 'node2')]
        cust,prod = edge_index_dict[('node2', 'to', 'node3')]
        attr1 = edge_attr_dict[('node1', 'to', 'node2')]
        attr2 = edge_attr_dict[('node2', 'to', 'node3')]
        # Concatenate the embeddings and edge attributed
        z1 = leaky_relu(z_dict['node1'][ship])
        z2 = leaky_relu(z_dict['node2'][cust])
        z3 = leaky_relu(z_dict['node3'][prod])
        z4 = leaky_relu(torch.cat([z1,z2,z3,attr1,attr2], dim=-1))
        z4 = self.network(z4)

        return z4.view(-1)


import setuptools

with open("README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()

setuptools.setup(
    name="transbigdata",
    version="0.1.0",
    author="Qing Yu",
    author_email="qingyu0815@foxmail.com",
    description="A tool for transportation big data",
    long_description=long_description,
    long_description_content_type="text/markdown",
    license="BSD",
    url="https://github.com/ni1o1/transbigdata",
    project_urls={
        "Bug Tracker": "https://github.com/ni1o1/transbigdata/issues",
    },
    install_requires=[
        "geopandas","matplotlib","plot_map","CoordinatesConverter"
        ],
    classifiers=[
        "Environment :: Web Environment",
        "Intended Audience :: Developers",
        "Operating System :: OS Independent",
        "Framework :: Matplotlib",
        "Topic :: Text Processing :: Indexing",
        "Topic :: Utilities",
        "Topic :: Internet",
        "Topic :: Software Development :: Libraries :: Python Modules",
        "Programming Language :: Python",
        "Programming Language :: Python :: 3.8",
    ],
    package_dir={"": "src"},
    packages=setuptools.find_packages(where="src"),
    python_requires=">=3.6",
)


from DB_Handler import DB_Handler
#program to safely store Ip addresses without hardcoding them into scripts
IPs = {}

DB = DB_Handler("Personal.db")

data = DB.RetrieveDB("IP_addresses")
for i in data:
    IPs[i[1]] = i[0]
    #print(i)

def getaddress(name):
    return(IPs[name])

text = input('')
text = text.split()
out = ""
out+= text[2]
out += " "
out += getaddress(text[0])
out += f":/{text[1]}"
#getaddress(name)
print(out)


def key_positions(seq, key):
    k = max(key(a) for a in seq)
    C = [0 for i in range(k+1)]
    for i in range(0, k):
        C[i] = 0
    for a in seq:
        C[key(a)] = (C[key(a)]) + 1
    sum = 0
    for i in range(0, k+1):
        C[i], sum = sum, sum+C[i]
    return C

#print(key_positions([0, 1, 2], lambda x: x))
#print(key_positions([2, 1, 0], lambda x: x))
#print(key_positions([1, 2, 3, 2], lambda x: x))
#print(key_positions([5], lambda x: x))
#print(key_positions(range(-3,3), lambda x: x**2))
#print(key_positions(range(1000), lambda x: 4))
#print(key_positions([1] + [0] * 100, lambda x: x))

def sorted_array(seq, key, positions):
    n = len(seq)
    B = [0]*n
    for a in seq:
        B[positions[key(a)]] = a
        positions[key(a)] = positions[key(a)] +1
    return B

#print(sorted_array([3, 1, 2], lambda x: x, [0, 0, 1, 2]))
#print(sorted_array([3, 2, 2, 1, 2], lambda x: x, [0, 0, 1, 4]))
#print(sorted_array([100], lambda x: x, [0]*101))
#"""Counting Sort"""
#import operator

#def counting_sort(iterable, key):
#    positions = key_positions(iterable, key)
#    return sorted_array(iterable, key, positions)

#objects = [("a", 88), ("b", 17), ("c", 17), ("d", 7)]

#key = operator.itemgetter(1)
#print(", ".join(object[0] for object in counting_sort(objects, key)))

def key_positions(seq, key):
    k = max(key(a) for a in seq)
    C = [0 for i in range(k+1)]
    for i in range(0, k):
        C[i] = 0
    for a in seq:
        C[key(a)] = (C[key(a)]) + 1
    sum = 0
    for i in range(0, k+1):
        C[i], sum = sum, sum+C[i]
    return C

def sorted_array(seq, key, positions):
    n = len(seq)
    B = [0]*n
    for a in seq:
        B[positions[key(a)]] = a
        positions[key(a)] = positions[key(a)] +1
    return B

def counting_sort(iterable, key):
    positions = key_positions(iterable, key)
    return sorted_array(iterable, key, positions)

#print(counting_sort([329, 457, 657, 839, 436, 720, 355], 3))
#print(radix_sort([329, 457, 657, 839, 436, 720, 355], 1))
#print(radix_sort([329, 457, 657, 839, 436, 720, 355], 2))

def sort(UnsortedList,position):
    bucket = []#this list will contain 10 list , each list is for every digit
    for i in range(0,10):
        poslist = []#list for integer i
        bucket.append(poslist)

    for num in UnsortedList:
        numstr = str(num)#converting number into string
        posval = 0
        if(len(numstr) >= position):#if number length is greater than equal to position
            #ord() function give ASCII value of the character
            #numstr[-position] will give the "position"th character from left
            posval = ord(numstr[-position]) - ord('0') #getting the value of "position" digit in every number in the list
            bucket[posval].append(num)#appending the number in the respective list [ if number is zero means the number will be appended into first list, if number is 1 then in 2nd list]
        else:
            bucket[0].append(num)

    sortedList = []#list will contain the numbers sorted according to the position

    #adding the numbers in sortedList from bucket
    for poslist in bucket:
        for num in poslist:
            sortedList.append(num)

    return sortedList

def radix_sort(numbers, d):
    finalList = numbers
    for position in range(1,d+1):#call sort function for position 1 to maxpos
        finalList = sort(finalList, position)

    return finalList

print(radix_sort([329, 457, 657, 839, 436, 720, 355], 3))
print(radix_sort([329, 457, 657, 839, 436, 720, 355], 1))
print(radix_sort([329, 457, 657, 839, 436, 720, 355], 2))
print(radix_sort([9, 57, 657], 1))
print(radix_sort([9, 57, 657], 2))
print(radix_sort([9, 57, 657], 3))


class Solution:
    def minCostII(self, costs: List[List[int]]) -> int:

        n = len(costs)
        k = len(costs[0])

        for house in range(1, n):
            min_color, second_min = None, None

            for color in range(k):
                cost = costs[house-1][color]
                if min_color is None or cost < costs[house - 1][min_color]:
                    second_min = min_color
                    min_color = color

                elif second_min is None or cost < costs[house - 1][second_min]:
                    second_min = color

            for color in range(k):
                if color == min_color:
                    costs[house][color] += costs[house - 1][second_min]
                else:
                    costs[house][color] += costs[house - 1][min_color]

        return (min(costs[-1]))


from django.urls import path
from . import views

app_name = "products"
urlpatterns = [
    path('', views.ProductList.as_view(), name="home"),
    path('new/', views.ProductCreate.as_view(), name="create"),
    path('<int:pk>/', views.ProductDetail.as_view(), name="detail"),
    path('<int:pk>/edit/', views.ProductEdit.as_view(), name="edit"),
    path('<int:pk>/delete/', views.ProductDelete.as_view(), name="delete"),
    path('csv/', views.generate_csv_view, name="generate_csv"),
]


# --------------------------------------------------------
# Adapted from  https://github.com/microsoft/unilm/tree/master/beit
# --------------------------------------------------------
import math
import os
from functools import partial
import logging
import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import drop_path, to_2tuple, trunc_normal_

#from .transformer import PatchDropout
from .layers.rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast

import numpy as np

if os.getenv('ENV_TYPE') == 'deepspeed':
    from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
else:
    from torch.utils.checkpoint import checkpoint
'''
try:
    import xformers.ops as xops
except ImportError:
    xops = None
    print("Please 'pip install xformers'")
'''

class PatchDropout(nn.Module):
    """
    https://arxiv.org/abs/2212.00794
    """

    def __init__(self, prob, exclude_first_token=True):
        super().__init__()
        assert 0 <= prob < 1.
        self.prob = prob
        self.exclude_first_token = exclude_first_token  # exclude CLS token
        logging.info(f"os.getenv('RoPE')={os.getenv('RoPE')}")

    def forward(self, x):
        if not self.training or self.prob == 0.:
            return x

        if self.exclude_first_token:
            cls_tokens, x = x[:, :1], x[:, 1:]
        else:
            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])

        batch = x.size()[0]
        num_tokens = x.size()[1]

        batch_indices = torch.arange(batch)
        batch_indices = batch_indices[..., None]

        keep_prob = 1 - self.prob
        num_patches_keep = max(1, int(num_tokens * keep_prob))

        rand = torch.randn(batch, num_tokens)
        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices

        x = x[batch_indices, patch_indices_keep]

        if self.exclude_first_token:
            x = torch.cat((cls_tokens, x), dim=1)

        if self.training and os.getenv('RoPE') == '1':
            return x, patch_indices_keep

        return x

class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob=None):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)

    def extra_repr(self) -> str:
        return 'p={}'.format(self.drop_prob)

class Mlp(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        drop=0.,
        subln=False,

        ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()

        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()

        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        # x = self.drop(x)
        # commit this for the orignal BERT implement
        x = self.ffn_ln(x)

        x = self.fc2(x)
        x = self.drop(x)
        return x

class SwiGLU(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
                norm_layer=nn.LayerNorm, subln=False):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features

        self.w1 = nn.Linear(in_features, hidden_features)
        self.w2 = nn.Linear(in_features, hidden_features)

        self.act = act_layer()
        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
        self.w3 = nn.Linear(hidden_features, out_features)

        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x1 = self.w1(x)
        x2 = self.w2(x)
        hidden = self.act(x1) * x2
        x = self.ffn_ln(hidden)
        x = self.w3(x)
        x = self.drop(x)
        return x

class SwiGLU_sep(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.,
                norm_layer=nn.LayerNorm, subln=False, num_det_tokens=100):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features

        self.w1_patch = nn.Linear(in_features, hidden_features)
        self.w1_det = nn.Linear(in_features, hidden_features)
        self.w2_patch = nn.Linear(in_features, hidden_features)
        self.w2_det = nn.Linear(in_features, hidden_features)

        self.act = act_layer()
        self.ffn_ln_patch = norm_layer(hidden_features) if subln else nn.Identity()
        self.ffn_ln_det = norm_layer(hidden_features) if subln else nn.Identity()
        self.w3_patch = nn.Linear(hidden_features, out_features)
        self.w3_det = nn.Linear(hidden_features, out_features)

        self.drop = nn.Dropout(drop)

        self.num_det_tokens = num_det_tokens

    def forward(self, x):
        x_patch = x[:,:-self.num_det_tokens,:]
        x_det = x[:,-self.num_det_tokens:,:]
        # x1 = self.w1(x)
        x1_patch = self.w1_patch(x_patch)
        x1_det = self.w1_det(x_det)
        # x2 = self.w2(x)
        x2_patch = self.w2_patch(x_patch)
        x2_det = self.w2_det(x_det)
        # hidden = self.act(x1) * x2
        hidden_patch = self.act(x1_patch) * x2_patch
        hidden_det = self.act(x1_det) * x2_det
        # x = self.ffn_ln(hidden)
        x_patch = self.ffn_ln_patch(hidden_patch)
        x_det = self.ffn_ln_det(hidden_det)
        # x = self.w3(x)
        x_patch = self.w3_patch(x_patch)
        x_det = self.w3_det(x_det)
        x = torch.cat((x_patch, x_det), dim=1)
        x = self.drop(x)
        return x

def memory_efficient_attention_pytorch(query, key, value, attn_bias=None, p=0., scale=None, attn_mask=None):
    # query     [batch, seq_len, n_head, head_dim]
    # key       [batch, seq_len, n_head, head_dim]
    # value     [batch, seq_len, n_head, head_dim]
    # attn_bias [batch, n_head, seq_len, seq_len]

    if scale is None:
        scale = 1 / query.shape[-1] ** 0.5

    # BLHC -> BHLC
    #query = query.transpose(1, 2)
    #key = key.transpose(1, 2)
    #value = value.transpose(1, 2)

    query = query * scale
    # BHLC @ BHCL -> BHLL
    attn = query @ key.transpose(-2, -1)
    if attn_bias is not None:
        attn = attn + attn_bias

    if attn_mask is not None:
        attn_mask = attn_mask.bool()
        attn = attn.masked_fill(~attn_mask, float("-inf"))

    attn = attn.softmax(-1)
    attn = F.dropout(attn, p)
    # BHLL @ BHLC -> BHLC
    out = attn @ value
    # BHLC -> BLHC
    out = out.transpose(1, 2)
    return out

class Attention(nn.Module):
    def __init__(
            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,proj_drop=0., window_size=None,
            attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm, num_det_tokens=100):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.num_det_tokens = num_det_tokens
        if attn_head_dim is not None:
            head_dim = attn_head_dim
        all_head_dim = head_dim * self.num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.subln = subln
        if self.subln:
            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
        else:
            self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)

        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias = None
            self.v_bias = None

        if window_size:
            self.window_size = window_size
            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
            self.relative_position_bias_table = nn.Parameter(
                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
            # cls to token & token 2 cls & cls to cls

            # get pair-wise relative position index for each token inside the window
            coords_h = torch.arange(window_size[0])
            coords_w = torch.arange(window_size[1])
            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
            relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
            relative_coords[:, :, 1] += window_size[1] - 1
            relative_coords[:, :, 0] *= 2 * window_size[1] - 1
            relative_position_index = \
                torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
            relative_position_index[0, 0:] = self.num_relative_distance - 3
            relative_position_index[0:, 0] = self.num_relative_distance - 2
            relative_position_index[0, 0] = self.num_relative_distance - 1

            self.register_buffer("relative_position_index", relative_position_index)
        else:
            self.window_size = None
            self.relative_position_bias_table = None
            self.relative_position_index = None

        self.attn_drop = nn.Dropout(attn_drop)
        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
        # self.proj = nn.Linear(all_head_dim, all_head_dim)
        self.proj = nn.Linear(all_head_dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.xattn = xattn
        self.xattn_drop = attn_drop

        self.rope = rope

    def forward(self, x, rel_pos_bias=None, attn_mask=None):
        B, N, C = x.shape
        if self.subln:
            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
            # np.savetxt('q.csv',q[0].cpu().detach().numpy(),fmt='%.4f',delimiter=',')
            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)

            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
        else:

            qkv_bias = None
            if self.q_bias is not None:
                qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))

            qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
            q, k, v = qkv[0], qkv[1], qkv[2]

        if self.rope:
            # slightly fast impl
            q_t = q[:, :, 1:-self.num_det_tokens, :]
            ro_q_t = self.rope(q_t)
            q = torch.cat((q[:, :, :1, :], ro_q_t, q[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

            k_t = k[:, :, 1:-self.num_det_tokens, :]
            ro_k_t = self.rope(k_t)
            k = torch.cat((k[:, :, :1, :], ro_k_t, k[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

        if self.xattn:
            #q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
            #k = k.permute(0, 2, 1, 3)
            #v = v.permute(0, 2, 1, 3)

            #x = xops.memory_efficient_attention(
            #    q, k, v,
            #    p=self.xattn_drop,
            #    scale=self.scale,
            #    )
            x = memory_efficient_attention_pytorch(q, k, v, p=self.xattn_drop, scale=self.scale, attn_mask=attn_mask)

            x = x.reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x = self.proj(x)
            x = self.proj_drop(x)
        else:
            q = q * self.scale
            attn = (q @ k.transpose(-2, -1))

            if self.relative_position_bias_table is not None:
                relative_position_bias = \
                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
                        self.window_size[0] * self.window_size[1] + 1,
                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)

            if rel_pos_bias is not None:
                attn = attn + rel_pos_bias.type_as(attn)

            if attn_mask is not None:
                attn_mask = attn_mask.bool()
                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))

            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)

            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x = self.proj(x)
            x = self.proj_drop(x)
        return x

class Attention_sep(nn.Module):
    def __init__(
            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,proj_drop=0., window_size=None,
            attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm, num_det_tokens=100):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.num_det_tokens = num_det_tokens
        if attn_head_dim is not None:
            head_dim = attn_head_dim
        all_head_dim = head_dim * self.num_heads
        self.scale = qk_scale or head_dim ** -0.5

        self.subln = subln
        if self.subln:
            self.q_proj_patch = nn.Linear(dim, all_head_dim, bias=False)
            self.k_proj_patch = nn.Linear(dim, all_head_dim, bias=False)
            self.v_proj_patch = nn.Linear(dim, all_head_dim, bias=False)
            self.q_proj_det = nn.Linear(dim, all_head_dim, bias=False)
            self.k_proj_det = nn.Linear(dim, all_head_dim, bias=False)
            self.v_proj_det = nn.Linear(dim, all_head_dim, bias=False)
        else:
            self.qkv_patch = nn.Linear(dim, all_head_dim * 3, bias=False)
            self.qkv_det = nn.Linear(dim, all_head_dim * 3, bias=False)

        if qkv_bias:
            self.q_bias_patch = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias_patch = nn.Parameter(torch.zeros(all_head_dim))
            self.q_bias_det = nn.Parameter(torch.zeros(all_head_dim))
            self.v_bias_det = nn.Parameter(torch.zeros(all_head_dim))
        else:
            self.q_bias_patch = None
            self.v_bias_patch = None
            self.q_bias_det = None
            self.v_bias_det = None

        if window_size:
            raise NotImplementedError
        else:
            self.window_size = None
            self.relative_position_bias_table = None
            self.relative_position_index = None

        self.attn_drop = nn.Dropout(attn_drop)
        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
        # self.proj = nn.Linear(all_head_dim, all_head_dim)
        self.proj_patch = nn.Linear(all_head_dim, dim)
        self.proj_det = nn.Linear(all_head_dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.xattn = xattn
        self.xattn_drop = attn_drop

        self.rope = rope

    def forward(self, x, rel_pos_bias=None, attn_mask=None):
        B, N, C = x.shape
        if self.subln:
            q_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.q_proj_patch.weight, bias=self.q_bias_patch)
            k_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.k_proj_patch.weight, bias=None)
            v_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.v_proj_patch.weight, bias=self.v_bias_patch)

            q_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.q_proj_det.weight, bias=self.q_bias_det)
            k_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.k_proj_det.weight, bias=None)
            v_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.v_proj_det.weight, bias=self.v_bias_det)

            q = torch.cat((q_patch, q_det), dim=1)
            k = torch.cat((k_patch, k_det), dim=1)
            v = torch.cat((v_patch, v_det), dim=1)

            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)
        else:
            qkv_bias_patch = None
            qkv_bias_det = None
            if self.q_bias_patch is not None:
                qkv_bias_patch = torch.cat((self.q_bias_patch, torch.zeros_like(self.v_bias_patch, requires_grad=False), self.v_bias_patch))
            if self.q_bias_det is not None:
                qkv_bias_det = torch.cat((self.q_bias_det, torch.zeros_like(self.v_bias_det, requires_grad=False), self.v_bias_det))

            qkv_patch = F.linear(input=x[:,:-self.num_det_tokens,:], weight=self.qkv_patch.weight, bias=qkv_bias_patch)
            qkv_det = F.linear(input=x[:,-self.num_det_tokens:,:], weight=self.qkv_det.weight, bias=qkv_bias_det)
            qkv = torch.cat((qkv_patch, qkv_det), dim=1)
            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
            q, k, v = qkv[0], qkv[1], qkv[2]

        if self.rope:
            # slightly fast impl
            q_t = q[:, :, 1:-self.num_det_tokens, :]
            ro_q_t = self.rope(q_t)
            q = torch.cat((q[:, :, :1, :], ro_q_t, q[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

            k_t = k[:, :, 1:-self.num_det_tokens, :]
            ro_k_t = self.rope(k_t)
            k = torch.cat((k[:, :, :1, :], ro_k_t, k[:, :, -self.num_det_tokens:, :]), -2).type_as(v)

        if self.xattn:
            #q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
            #k = k.permute(0, 2, 1, 3)
            #v = v.permute(0, 2, 1, 3)

            #x = xops.memory_efficient_attention(
            #    q, k, v,
            #    p=self.xattn_drop,
            #    scale=self.scale,
            #    )
            x = memory_efficient_attention_pytorch(q, k, v, p=self.xattn_drop, scale=self.scale, attn_mask=attn_mask)

            x = x.reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x_patch = self.proj_patch(x[:,:-self.num_det_tokens,:])
            x_det = self.proj_det(x[:,-self.num_det_tokens:,:])
            x = torch.cat((x_patch, x_det), dim=1)
            x = self.proj_drop(x)
        else:
            q = q * self.scale
            attn = (q @ k.transpose(-2, -1))

            if self.relative_position_bias_table is not None:
                relative_position_bias = \
                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
                        self.window_size[0] * self.window_size[1] + 1,
                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)

            if rel_pos_bias is not None:
                attn = attn + rel_pos_bias.type_as(attn)

            if attn_mask is not None:
                attn_mask = attn_mask.bool()
                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))

            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)

            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
            x = self.inner_attn_ln(x)
            x = self.proj(x)
            x = self.proj_drop(x)
        return x

class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
                 window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
                 subln=False, naiveswiglu=False, num_det_tokens=100, partial_finetune=False, partial_finetune_attn=False):
        super().__init__()
        self.norm1 = norm_layer(dim)
        if partial_finetune_attn:
            self.attn = Attention_sep(
                dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
                xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer, num_det_tokens=num_det_tokens)
        else:
            self.attn = Attention(
                dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
                attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
                xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer, num_det_tokens=num_det_tokens)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)

        if naiveswiglu:
            if partial_finetune:
                self.mlp = SwiGLU_sep(
                    in_features=dim,
                    hidden_features=mlp_hidden_dim,
                    subln=subln,
                    norm_layer=norm_layer,
                    num_det_tokens=num_det_tokens,
                )
            else:
                self.mlp = SwiGLU(
                    in_features=dim,
                    hidden_features=mlp_hidden_dim,
                    subln=subln,
                    norm_layer=norm_layer,
                )
        else:
            self.mlp = Mlp(
                in_features=dim,
                hidden_features=mlp_hidden_dim,
                act_layer=act_layer,
                subln=subln,
                drop=drop
            )

        if init_values is not None and init_values > 0:
            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
        else:
            self.gamma_1, self.gamma_2 = None, None

        self.postnorm = postnorm

    def forward(self, x, rel_pos_bias=None, attn_mask=None):
        if self.gamma_1 is None:
            if self.postnorm:
                x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
                x = x + self.drop_path(self.norm2(self.mlp(x)))
            else:
                x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
                x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            if self.postnorm:
                x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
            else:
                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
        return x

class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

    def forward(self, x, **kwargs):
        B, C, H, W = x.shape
        # FIXME look at relaxing size constraints
        #assert H == self.img_size[0] and W == self.img_size[1], \
        #    f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x).flatten(2).transpose(1, 2)  # [1, 14*14, vision_cfg.width:768]
        return x

class RelativePositionBias(nn.Module):

    def __init__(self, window_size, num_heads):
        super().__init__()
        self.window_size = window_size
        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
        # cls to token & token 2 cls & cls to cls

        # get pair-wise relative position index for each token inside the window
        coords_h = torch.arange(window_size[0])
        coords_w = torch.arange(window_size[1])
        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
        relative_coords[:, :, 1] += window_size[1] - 1
        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
        relative_position_index = \
            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
        relative_position_index[0, 0:] = self.num_relative_distance - 3
        relative_position_index[0:, 0] = self.num_relative_distance - 2
        relative_position_index[0, 0] = self.num_relative_distance - 1

        self.register_buffer("relative_position_index", relative_position_index)

    def forward(self):
        relative_position_bias = \
            self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
                self.window_size[0] * self.window_size[1] + 1,
                self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww

class EVAVisionTransformer(nn.Module):
    """ Vision Transformer with support for patch or hybrid CNN input stage
    """
    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, num_det_tokens=100, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,drop_path_rate=0.,
                 norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,use_abs_pos_emb=True, use_rel_pos_bias=False,
                 use_shared_rel_pos_bias=False, rope=False,use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False,
                 xattn=False, postnorm=False,pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False, attn_mask=False,
                 partial_finetune=False, partial_finetune_attn=False):
        super().__init__()
        #self.img_size = img_size
        if isinstance(img_size,tuple):
            self.img_size = img_size
        else:
            self.img_size = to_2tuple(img_size)
        self.num_classes = num_classes
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models

        self.num_det_tokens = num_det_tokens

        self.patch_embed = PatchEmbed(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
        self.patch_size = patch_size
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        # self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        if use_abs_pos_emb:
            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        else:
            self.pos_embed = None
        self.pos_drop = nn.Dropout(p=drop_rate)

        if use_shared_rel_pos_bias:
            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
        else:
            self.rel_pos_bias = None

        if rope:
            half_head_dim = embed_dim // num_heads // 2
            hw_seq_len = img_size // patch_size
            self.rope = VisionRotaryEmbeddingFast(
                dim=half_head_dim,
                pt_seq_len=pt_hw_seq_len,
                ft_seq_len=hw_seq_len if intp_freq else None,
                # patch_dropout=patch_dropout
            )
        else:
            self.rope = None

        self.naiveswiglu = naiveswiglu
        self.partial_finetune = partial_finetune
        self.partial_finetune_attn = partial_finetune_attn

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        self.use_rel_pos_bias = use_rel_pos_bias
        self.blocks = nn.ModuleList([
            Block(
                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
                xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu,
                num_det_tokens = self.num_det_tokens, partial_finetune=partial_finetune, partial_finetune_attn = partial_finetune_attn)
            for i in range(depth)])
        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
        self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()

        if self.pos_embed is not None:
            trunc_normal_(self.pos_embed, std=.02)

        trunc_normal_(self.cls_token, std=.02)
        # trunc_normal_(self.mask_token, std=.02)

        self.apply(self._init_weights)
        self.fix_init_weight()

        if isinstance(self.head, nn.Linear):
            trunc_normal_(self.head.weight, std=.02)
            self.head.weight.data.mul_(init_scale)
            self.head.bias.data.mul_(init_scale)

        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()

        self.grad_checkpointing = grad_checkpointing
        self.has_mid_pe = False

        self.attn_mask = attn_mask

    def fix_init_weight(self):
        def rescale(param, layer_id):
            param.div_(math.sqrt(2.0 * layer_id))

        for layer_id, layer in enumerate(self.blocks):
            if self.partial_finetune_attn:
                rescale(layer.attn.proj_patch.weight.data, layer_id + 1)
                rescale(layer.attn.proj_det.weight.data, layer_id + 1)
            else:
                rescale(layer.attn.proj.weight.data, layer_id + 1)
            if self.naiveswiglu:
                if self.partial_finetune:
                    rescale(layer.mlp.w3_patch.weight.data, layer_id + 1)
                    rescale(layer.mlp.w3_det.weight.data, layer_id + 1)
                else:
                    rescale(layer.mlp.w3.weight.data, layer_id + 1)
            else:
                rescale(layer.mlp.fc2.weight.data, layer_id + 1)

    def get_cast_dtype(self) -> torch.dtype:
        return self.blocks[0].mlp.fc2.weight.dtype

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def get_num_layers(self):
        return len(self.blocks)

    def lock(self, unlocked_groups=0):
        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
        for param in self.parameters():
            param.requires_grad = False

    @torch.jit.ignore
    def set_grad_checkpointing(self, enable=True):
        self.grad_checkpointing = enable

    @torch.jit.ignore
    def no_weight_decay(self):
        return {'pos_embed', 'cls_token'}

    def get_classifier(self):
        return self.head

    def reset_classifier(self, num_classes, global_pool=''):
        self.num_classes = num_classes
        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()

    def finetune_det(self, img_size=[800, 1344], mid_pe_size=None, use_checkpoint=False):
        # import pdb;pdb.set_trace()

        import math
        g = math.pow(self.pos_embed.size(1) - 1, 0.5)
        if int(g) - g != 0:
            self.pos_embed = torch.nn.Parameter(self.pos_embed[:, 1:, :])

        # self.det_token_num = det_token_num
        self.det_token = nn.Parameter(torch.zeros(1, self.num_det_tokens, self.embed_dim))
        self.det_token = trunc_normal_(self.det_token, std=.02)
        cls_pos_embed = self.pos_embed[:, 0, :]
        cls_pos_embed = cls_pos_embed[:,None]
        det_pos_embed = torch.zeros(1, self.num_det_tokens, self.embed_dim)
        det_pos_embed = trunc_normal_(det_pos_embed, std=.02)
        patch_pos_embed = self.pos_embed[:, 1:, :]
        patch_pos_embed = patch_pos_embed.transpose(1,2)
        B, E, Q = patch_pos_embed.shape
        P_H, P_W = self.img_size[0] // self.patch_size, self.img_size[1] // self.patch_size
        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)
        H, W = img_size
        new_P_H, new_P_W = H//self.patch_size, W//self.patch_size
        # torch.Size([1, 768, 50, 84])
        patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        patch_pos_embed = patch_pos_embed.flatten(2).transpose(1,2) # torch.Size([1, 4200, 768])
        self.pos_embed = torch.nn.Parameter(torch.cat((cls_pos_embed, patch_pos_embed), dim=1), requires_grad=not (self.partial_finetune or self.partial_finetune_attn))
        self.det_pos_embed = torch.nn.Parameter(det_pos_embed)

        _, head_dim = self.rope.freqs_cos.shape
        freqs_cos = self.rope.freqs_cos.transpose(0,1).view(B, head_dim, P_H, P_W)
        freqs_sin = self.rope.freqs_sin.transpose(0,1).view(B, head_dim, P_H, P_W)
        freqs_cos = nn.functional.interpolate(freqs_cos, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False) #torch.Size([1, 64, 50, 84])
        freqs_sin = nn.functional.interpolate(freqs_sin, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        self.register_buffer("freqs_cos", freqs_cos.squeeze().flatten(1).transpose(0,1))
        self.register_buffer("freqs_sin", freqs_sin.squeeze().flatten(1).transpose(0,1))

        self.img_size = img_size
        if mid_pe_size == None:
            self.has_mid_pe = False
            print('No mid pe')
        else:
            print('Has mid pe')
            self.mid_pos_embed = nn.Parameter(torch.zeros(self.depth - 1, 1, 1 + (mid_pe_size[0] * mid_pe_size[1] // self.patch_size ** 2) + self.num_det_tokens, self.embed_dim))
            trunc_normal_(self.mid_pos_embed, std=.02)
            self.has_mid_pe = True
            self.mid_pe_size = mid_pe_size
        self.use_checkpoint=use_checkpoint

    def InterpolateInitPosEmbed(self, pos_embed, img_size=(800, 1344)):
        # import pdb;pdb.set_trace()
        _, head_dim = self.rope.freqs_cos.shape

        cls_pos_embed = pos_embed[:, 0, :]
        cls_pos_embed = cls_pos_embed[:,None]
        patch_pos_embed = pos_embed[:, 1:, :] # [1,4200,768]
        patch_pos_embed = patch_pos_embed.transpose(1,2)
        freqs_cos = self.freqs_cos.transpose(0,1)
        freqs_sin = self.freqs_sin.transpose(0,1)
        B, E, Q = patch_pos_embed.shape

        P_H, P_W = self.img_size[0] // self.patch_size, self.img_size[1] // self.patch_size
        patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)
        freqs_cos = freqs_cos.view(B, head_dim, P_H, P_W)
        freqs_sin = freqs_sin.view(B, head_dim, P_H, P_W)

        # P_H, P_W = self.img_size[0] // self.patch_size, self.img_size[1] // self.patch_size
        # patch_pos_embed = patch_pos_embed.view(B, E, P_H, P_W)

        H, W = img_size
        new_P_H, new_P_W = H//self.patch_size, W//self.patch_size
        patch_pos_embed = nn.functional.interpolate(patch_pos_embed, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        patch_pos_embed = patch_pos_embed.flatten(2).transpose(1, 2)
        scale_pos_embed = torch.cat((cls_pos_embed, patch_pos_embed), dim=1)

        freqs_cos = nn.functional.interpolate(freqs_cos, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        freqs_sin = nn.functional.interpolate(freqs_sin, size=(new_P_H,new_P_W), mode='bicubic', align_corners=False)
        self.rope.freqs_cos = freqs_cos.squeeze().flatten(1).transpose(0,1)
        self.rope.freqs_sin = freqs_sin.squeeze().flatten(1).transpose(0,1)

        return scale_pos_embed

    def get_rect_mask(self, patch_seq_len=196):

        total_seq_len = 1 + patch_seq_len + self.num_det_tokens
        mask = torch.ones(size=[total_seq_len, total_seq_len], device='cuda')
        mask[:-self.num_det_tokens, -self.num_det_tokens:] = 0

        return mask

    def forward_features(self, x, return_all_features=False):

        B, H, W = x.shape[0], x.shape[2], x.shape[3]

        x = self.patch_embed(x)  # [B,3,H,W] -> torch.Size([B, n_patch, embed_dim=768])
        batch_size, seq_len, _ = x.size()

        # interpolate init pe
        if (self.pos_embed.shape[1] - 1) != seq_len:
            temp_pos_embed = self.InterpolateInitPosEmbed(self.pos_embed, img_size=(H,W))
        else:
            temp_pos_embed = self.pos_embed
            self.rope.freqs_cos = self.freqs_cos
            self.rope.freqs_sin = self.freqs_sin

        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        det_token = self.det_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, x, det_token), dim=1)  # [B,1+n_patch+100, embed_dim=768]
        if self.pos_embed is not None:
            x = x + torch.cat((temp_pos_embed, self.det_pos_embed), dim=1)
        x = self.pos_drop(x)

        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
        if os.getenv('RoPE') == '1':
            if self.training and not isinstance(self.patch_dropout, nn.Identity):
                x, patch_indices_keep = self.patch_dropout(x)
                self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
            else:
                self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
                x = self.patch_dropout(x)  # Identity()
        else:
            x = self.patch_dropout(x)

        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None  # None
        attn_mask = self.get_rect_mask(patch_seq_len=seq_len) if self.attn_mask else None
        for blk in self.blocks:
            if self.grad_checkpointing:
                x = checkpoint(blk, x, (rel_pos_bias, attn_mask))
            else:
                x = blk(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)

        if not return_all_features:  # x:[B, 1+n_patch+100, embed_dim=768]
            x = self.norm(x)
            if self.fc_norm is not None:
                return self.fc_norm(x.mean(1))
            else:
                return x[:, -self.num_det_tokens:, :]  # torch.Size([1, 100, 768])
        return x

    def forward(self, x, return_all_features=False):
        if return_all_features:
            return self.forward_features(x, return_all_features)
        x = self.forward_features(x)
        # x = self.head(x) # [1,512]
        return x