# vim: fdm=indent
'''
author:     Fabio Zanini
date:       09/01/23
content:    Explore first two samples, P7 normals.
'''
import os
import sys
import pathlib
import numpy as np
import pandas as pd
import anndata
import scanpy as sc

import matplotlib.pyplot as plt
import seaborn as sns

figures_fdn = pathlib.Path('../../figures/pilots/lungmo_Dec2022')

def plot_clustermap_avg(ge_avg, log=True):
    from scipy.spatial.distance import pdist
    from scipy.cluster.hierarchy import linkage, leaves_list

    if log:
        ge_avg = np.log(ge_avg + 0.1)

    pdis = pdist(ge_avg.values)
    Zrow = linkage(pdis, method='average', optimal_ordering=True)

    pdis = pdist(ge_avg.values.T, metric='correlation')
    Zcol = linkage(pdis, method='average', optimal_ordering=True)

    g = sns.clustermap(
        ge_avg, row_linkage=Zrow, col_linkage=Zcol,
        xticklabels=True, yticklabels=True,
    )

    return g

if __name__ == '__main__':

    print('Read h5ad file')
    data_fdn = pathlib.Path('../../data/pilot_Dec2022/counts')
    adata = anndata.read_h5ad(data_fdn / 'merged_nor-1_nor-2.h5ad')

    # NOTE: binarize ATAC
    #ngenes = (adata.var['feature_types'] == 'Gene Expression').sum()

    adata.obs['coverage_ge'] = adata[:, adata.var['feature_types'] == 'Gene Expression'].X.sum(axis=1)
    adata.obs['coverage_pe'] = adata[:, adata.var['feature_types'] == 'Peaks'].X.sum(axis=1)

    adatag = adata[adata.obs['coverage_ge'] >= 500]

    if False:
        print('Check heatmap of marker genes')
        marker_exp = pd.DataFrame([], index=adata.obs_names)
        genes = ['Cdh5', 'Pecam1',
                 'Ptprc',
                 'Epcam', 'Krt19',
                 'Col1a1', 'Col1a2', 'Col6a2',
                 'Snap25',
                 'Gja5', 'Bmx', 'Vwf', 'Car4', 'Peg3', 'Car8',
                 'Fn1', 'Ctsh', 'Kcne3', 'Cdh13', 'Thy1', 'Ccl21a',
                 'Sirpa', 'Fibin',
                 'Cd3d', 'Cd3e', 'Cd19', 'Ms4a1', 'Cd68', 'Plac8', 'Dab2', 'C1qa',
                 'Mcpt4', 'Mcpt8', 'Itgax', 'Retnlg', 'Gzma', 'Areg',
                 'Acta2', 'Pdgfrb', 'Hhip', 'Pdgfra', 'Adh1',
                 ]
        adatam = adatag[:, genes]

        X = adatam.X.todense()
        X = X / (X.sum(axis=1) + 0.1)

        from scipy.spatial.distance import pdist
        from scipy.cluster.hierarchy import linkage, leaves_list

        pdis = pdist(X)
        Z = linkage(pdis, method='average')
        idx_cells = leaves_list(Z)

        pdis = pdist(X.T, metric='correlation')
        Z = linkage(pdis, method='average')
        idx_genes = leaves_list(Z)
        gene_names = [genes[i] for i in idx_genes]

        Xcl = X[idx_cells].T[idx_genes].T

        fig, ax = plt.subplots()
        sns.heatmap(np.log(Xcl + 0.5), yticklabels=False, ax=ax, xticklabels=True)
        ax.set_xticklabels(gene_names, rotation=90)
        fig.tight_layout()

    if True:
        print('Some QC')
        adata_ge = adata[:, adata.var['feature_types'] == 'Gene Expression']
        adata_pe = adata[:, adata.var['feature_types'] == 'Peaks']
        sc.pp.normalize_total(adata_ge, target_sum=1e4)
        sc.pp.normalize_total(adata_pe, target_sum=1e4)
        adata_nor = anndata.concat([adata_ge, adata_pe], axis=1, merge='first')

        fig, ax = plt.subplots()
        colors = {'nor-1': 'tomato', 'nor-2': 'steelblue'}
        for sample, group in adata_nor.obs.groupby('sample'):
            ax.scatter(
                0.1 + group['coverage_ge'],
                0.1 + group['coverage_pe'],
                color=colors[sample],
                alpha=0.2,
                label=sample,
            )
        ax.set_xscale('log')
        ax.set_yscale('log')
        ax.set_xlabel('Coverage GEX [UMI]')
        ax.set_ylabel('Coverage ATAC [UMI]')
        ax.axvline(500, ls='--', lw=2, color='k')
        ax.grid(True)
        ax.legend()
        fig.tight_layout()

        fig.savefig(
            figures_fdn / 'QC_number_of_UMI_per_cell.png'
        )

    if False:
        print('Try standard clustering, with features from both')
        adatag = adata_nor[adata_nor.obs['coverage_ge'] >= 500]
        adata_ge = adatag[:, adatag.var['feature_types'] == 'Gene Expression']

        # Log
        sc.pp.log1p(adata_ge)

        # HVG
        sc.pp.highly_variable_genes(adata_ge, min_mean=0.0125, max_mean=3, min_disp=0.5)

        adata_ge.raw = adata_ge
        adata_ge = adata_ge[:, adata_ge.var.highly_variable]

        # scale
        sc.pp.scale(adata_ge, max_value=10)

        # PCA
        sc.tl.pca(adata_ge, svd_solver='arpack')

        # neighborhood graph
        sc.pp.neighbors(adata_ge, n_neighbors=10, n_pcs=40)

        # Embed
        sc.tl.umap(adata_ge)

        # Cluster
        sc.tl.leiden(adata_ge)

        # Compute avg expression of markers
        genes = ['Cdh5', 'Pecam1',
                 'Ptprc',
                 'Epcam', 'Krt19',
                 'Col1a1', 'Col1a2', 'Col6a2',
                 'Snap25',
                 'Gja5', 'Bmx', 'Vwf', 'Car4', 'Peg3', 'Car8',
                 'Fn1', 'Ctsh', 'Kcne3', 'Cdh13', 'Thy1', 'Ccl21a',
                 'Sirpa', 'Fibin',
                 'Cd3d', 'Cd3e', 'Cd19', 'Ms4a1', 'Cd68', 'Plac8', 'Dab2', 'C1qa',
                 'Mcpt4', 'Mcpt8', 'Itgax', 'Retnlg', 'Gzma', 'Areg',
                 'Acta2', 'Pdgfrb', 'Hhip', 'Pdgfra', 'Adh1',
                 ]
        if adata_ge.raw is not None:
            adata_tmp = adata_ge.raw.to_adata()[:, genes]
        else:
            adata_tmp = adata_ge[:, genes]
        clusters = adata_ge.obs['leiden'].cat.categories
        ge_avg = np.zeros((len(clusters), len(genes)))
        for i, clu in enumerate(clusters):
            adatai = adata_tmp[adata_tmp.obs['leiden'] == clu]
            avgi = adatai.X.mean(axis=0)
            ge_avg[i] = avgi
        ge_avg = pd.DataFrame(ge_avg, index=clusters, columns=genes)

        plot_clustermap_avg(ge_avg, log=False)

    if True:
        print('Assign cell types to clusters')
        adata_ge = adata[:, adata.var['feature_types'] == 'Gene Expression']
        adata_pe = adata[:, adata.var['feature_types'] == 'Peaks']
        sc.pp.normalize_total(adata_ge, target_sum=1e4)
        sc.pp.normalize_total(adata_pe, target_sum=1e4)
        adata_nor = anndata.concat([adata_ge, adata_pe], axis=1, merge='first')

        assignment_d = {
            'Gja5': 'Arterial EC',
            'Ccl21a': 'Lymphatic EC',
            'Slc6a2': 'Venous EC',
            'Hpgd': 'CAP1',
            'Car4': 'CAP2',
            'C1qa': 'Interstitial mac',
            'Itgax': 'Alveolar mac',
            'Plac8': 'Monocyte',
            'Itgae': 'cDC1',
            'Cd209a': 'cDC2',
            #'Mreg': 'cDC3',
            'Mcpt8': 'Mast',
            'Mcpt4': 'Basophil',
            'Retnlg': 'Neutrophil',
            'Ms4a1': 'B cell',
            'Gzma': 'NK cell',
            'Cd3e': 'T cell',
            'Areg': 'ILC2',
            'Col13a1': 'Alveolar FB',
            'Col14a1': 'Adventitial FB',
            'Hhip': 'ASM',
            'Higd1b': 'Pericyte',
            'Prrx1': 'VSM',
            'Actc1': 'ASM',
            'Slc34a2': 'AT2',
            'Col4a4': 'AT1',
            'Cdkn1c': 'Ciliated',
            'Reg3g': 'Club',
            'Snap25': 'Neuron',
        }
        other_genes = [
                'Ptprc', 'Cdh5', 'Epcam', 'Col6a2',
                'Cd68', 'Pecam1', 'Krt19', 'Col1a1', 'Lyz1',
                'Lyz2', 'Pdgfra', 'Pdgfrb', 'Tgfbi',
                ]
        genes = list(assignment_d.keys()) + other_genes
        if adata_ge.raw is not None:
            adata_tmp = adata_ge.raw.to_adata()[:, genes]
        else:
            adata_tmp = adata_ge[:, genes]
        clusters = adata_ge.obs['leiden'].cat.categories
        ge_avg = np.zeros((len(clusters), len(genes)))
        for i, clu in enumerate(clusters):
            adatai = adata_tmp[adata_tmp.obs['leiden'] == clu]
            avgi = adatai.X.mean(axis=0)
            ge_avg[i] = avgi
        ge_avg = pd.DataFrame(ge_avg, index=clusters, columns=genes)

        # Automatic detection based on marker genes
        cell_type_d = {}
        for gene in assignment_d:
            cell_type_d[ge_avg[gene].idxmax()] = assignment_d[gene]
        for clu in clusters:
            if clu not in cell_type_d:
                cell_type_d[clu] = clu

        ge_avg_tmp = ge_avg.copy()
        ge_avg_tmp.index = ge_avg_tmp.index.map(cell_type_d)
        #plot_clustermap_avg(ge_avg_tmp, log=True)

        # Manual corrections
        manual_corr = {
            '0': 'Alveolar FB',
            '1': 'Low-quality',
            '2': 'CAP1',
            '3': 'Epithelial',
            '4': 'Myofibroblast',
            '5': 'Adventitial FB',
            '7': 'Alveolar FB',
            '8': 'Club',
            '11': 'Doublet',
            '17': 'Unknown',
            '20': 'VSM',
            '23': 'Doublet',
            '25': 'Venous EC',
            '29': 'Doublet',
            '30': 'Doublet',
            '32': 'Doublet',

        }
        cell_type_d.update(manual_corr)

        ge_avg_tmp = ge_avg.copy()
        ge_avg_tmp.index = ge_avg_tmp.index.map(cell_type_d)
        plot_clustermap_avg(ge_avg_tmp, log=True)

        adata_ge.obs['Cell Type'] = pd.Categorical(
            adata_ge.obs['leiden'].map(cell_type_d),
        )
        adata_ge.obs.loc[adata_ge.obs['coverage_ge'] < 500, 'Cell Type'] = 'Low-quality'
        adata_ge.obs['High-quality'] = (~adata_ge.obs['Cell Type'].isin(
                ['Doublet', 'Low-quality', 'Unknown'],
                )).astype('i2')
        adata_nor.obs['Cell Type'] = adata_ge.obs['Cell Type']
        adata_nor.obs['High-quality'] = adata_ge.obs['High-quality']

    if True:
        print('Report how many cells per type/subtype we got')
        df = adata_nor.obs.loc[adata_nor.obs['High-quality'] == 1].copy()
        df['Cell Type'] = df['Cell Type'].cat.remove_unused_categories()
        cst_abu = df.groupby(['Cell Type', 'sample']).size().unstack(fill_value=0)
        cst_abu.to_csv(
            figures_fdn / 'number_of_cells_per_sample_and_subtype.tsv',
            sep='\t', index=True,
        )

    if True:
        print('Show umap of annotated subtypes')
        adatag = adata_nor[adata_nor.obs['High-quality'] == 1]
        adatag.obs['Cell Type'] = adatag.obs['Cell Type'].cat.remove_unused_categories()

        df = pd.read_csv(
            '../../data/pilot_Dec2022/lung_multiomics_leiden_and_UMAP.csv',
            index_col=0,
        ).loc[adatag.obs_names]
        adatag.obsm['X_umap'] = df[['UMAP1', 'UMAP2']].values

        fig, ax = plt.subplots(figsize=(8, 4))
        cell_types = adatag.obs['Cell Type'].unique()
        cmap = dict(zip(cell_types, sns.color_palette('husl', n_colors=len(cell_types))))
        df['Cell Type'] = adatag.obs['Cell Type']
        for i, (cst, grp) in enumerate(df.groupby('Cell Type')):
            x = grp['UMAP1']
            y = grp['UMAP2']
            marker = 'o' if i < 12 else 's'
            ax.scatter(
                x, y, color=cmap[cst], label=str(i+1) + '. ' + cst, alpha=0.3, s=30,
                marker=marker,
            )
            ax.text(x.mean(), y.mean(), str(i+1), ha='center', va='center')
        ax.legend(
            loc='center left', title='Pilot multiomics - Dec 2022',
            bbox_to_anchor=(1, 0.5), bbox_transform=ax.transAxes,
            ncol=2, frameon=False,
        )
        fig.tight_layout()
        fig.savefig(
            figures_fdn / 'pilot_UMAP_by_cell_type.png',
        )

    sys.exit()

    if False:
        print('DE and DA within lineages')
        lineages = {
            'Endothelial': [
                'Arterial EC',
                'Venous EC',
                'Lymphatic EC',
                'CAP1',
                'CAP2',
            ],
            'Immune': [
                'T cell',
                'B cell',
                'NK cell',
                'Alveolar mac',
                'Interstitial mac',
                'Monocyte',
            ],
            # TODO
        }

        lineage = 'Endothelial'
        adata_nor = adata_nor[adata_nor.obs['High-quality'] == 1]
        adata_nor_gr = adata_nor[adata_nor.obs['Cell Type'].isin(lineages[lineage])]
        adata_ge = adata_nor_gr[:, adata_nor_gr.var['feature_types'] == 'Gene Expression']
        adata_pe = adata_nor_gr[:, adata_nor_gr.var['feature_types'] == 'Peaks']
        sc.tl.rank_genes_groups(adata_pe, 'Cell Type', method='wilcoxon')
        sc.tl.rank_genes_groups(adata_ge, 'Cell Type', method='wilcoxon')

        from collections import defaultdict
        da_lineage_dict = defaultdict(list)
        for ftype, adatai in [('Peaks', adata_pe), ('Gene Expression', adata_ge)]:
            tmp = adatai.uns['rank_genes_groups']
            cell_types = tmp['names'].dtype.names
            for i, ct in enumerate(cell_types):
                idx = [x[i] for x in tmp['names']]
                log_fc = [x[i] for x in tmp['logfoldchanges']]
                score = [x[i] for x in tmp['scores']]
                pval = [x[i] for x in tmp['pvals']]
                df = pd.DataFrame({
                    'feature': idx,
                    'logfc': log_fc,
                    'score': score,
                    'pval': pval,
                }).set_index('feature')
                df.sort_values('score', ascending=False, inplace=True)
                df['rank'] = np.arange(len(df)) + 1

                for col in ['chromosome', 'start', 'end']:
                    df[col] = adatai.var.loc[df.index, col]
                df['feature_types'] = ftype
                da_lineage_dict[ct].append(df)
        for ct, dfs in da_lineage_dict.items():
            da_lineage_dict[ct] = pd.concat(dfs, axis=0)

    def plot_de_on_chroms(var, df):
        chromosomes = [f'chr{i+1}' for i in range(19)] + ['chrX', 'chrY']
        nchr = len(chromosomes)
        chr_max = (var[['chromosome', 'start', 'end']]
                      .groupby('chromosome')
                      .max()
                      .max(axis=1)
                      .loc[chromosomes]) + 1

        fig, ax = plt.subplots(figsize=(10, 10))

        # Chromosome tracks
        for i, chrom in enumerate(chromosomes):
            ax.plot([0, 100], [nchr - 1 - i] * 2, lw=1, color='k', alpha=0.3)

        cmap = {'Gene Expression': 'tomato', 'Peaks': 'steelblue'}
        offsets = {'Gene Expression': +0.1, 'Peaks': -0.1}
        for feature, row in df.iterrows():
            chrom = row['chromosome']
            if chrom not in chromosomes:
                continue
            ftype = row['feature_types']
            irow = nchr - 1 - chromosomes.index(chrom)
            x0 = 100.0 * row['start'] / chr_max[chrom]
            x1 = 100.0 * row['end'] / chr_max[chrom]
            # Min width 0.5% of chromosome, otherwise you don't even see them
            if (x1 - x0) < 0.5:
                x0 -= 0.25
                x1 += 0.25
            ax.plot(
                [x0, x1], [irow + offsets[ftype]] * 2,
                color=cmap[ftype], lw=3, alpha=0.8,
            )
            xm = 0.5 * (x0 + x1)
            if ftype == 'Gene Expression':
                ax.text(xm, irow + 0.15, feature, ha='center', va='bottom')
                strand = var.at[feature, 'strand']
                # Mark the beginning
                if strand == -1:
                    ax.plot(
                        [x0 + 0.8 * (x1 - x0), x1],
                        [irow + offsets[ftype]] * 2,
                        color=cmap[ftype], lw=4, alpha=0.9,
                        )
                else:
                    ax.plot(
                        [x0, x0 + 0.2 * (x1 - x0)],
                        [irow + offsets[ftype]] * 2,
                        color=cmap[ftype], lw=4, alpha=0.9,
                        )

        labels = ['Gene', 'Open chromatin']
        handles = [
            plt.Rectangle((0, 0), 0, 0, color=cmap['Gene Expression']),
            plt.Rectangle((0, 0), 0, 0, color=cmap['Peaks']),
        ]
        ax.legend(handles, labels,
                  loc='center',
                  bbox_to_anchor=(0.85, 1.015), bbox_transform=ax.transAxes,
                  ncol=2,
                  frameon=False,
                  )

        ax.set_yticks(np.arange(0, len(chromosomes)))
        ax.set_yticklabels(chromosomes[::-1])
        ax.set_xticks([0, 50, 100])
        ax.set_xlabel('Position within chromosome[%]')
        ax.set_ylim(-0.5, len(chromosomes) - 0.5)
        ax.set_xlim(-1, 101)
        fig.tight_layout()

        return (fig, ax)

    if False:
        for ct in da_lineage_dict:
            fig, ax = plot_de_on_chroms(
                adata.var,
                da_lineage_dict[ct].nsmallest(100, 'rank'),
            )
            ax.set_title(ct)
            fig.savefig(
                figures_fdn / '..' / 'de_da_across_chromosomes_endos' / f'{ct}.png',
            )

    def get_closest_features(var, feature, target_type):
        chrom, start, end = var.loc[feature][['chromosome', 'start', 'end']].values
        other = var.loc[(var['chromosome'] == chrom) & (var['feature_types'] == target_type)].copy()
        if len(other) == 0:
            return pd.Series([])

        is_overlap = ~((other['start'] > end) | (other['end'] < start))
        other_ov = other.loc[is_overlap].index
        if len(other_ov):
            return pd.Series(np.zeros(len(other_ov), int), index=other_ov)

        # No overlap, they are either before or after
        after = other.loc[other['start'] > end].copy()
        after['distance'] = after['start'] - end
        after_best = after['distance'].idxmin()
        before = other.loc[other['end'] < start].copy()
        before['distance'] = start - before['end']
        before_best = before['distance'].idxmin()

        if after.at[after_best, 'distance'] < before.at[before_best, 'distance']:
            return pd.Series([after.at[after_best, 'distance']], index=[after_best])
        elif after.at[after_best, 'distance'] > before.at[before_best, 'distance']:
            return pd.Series([before.at[before_best, 'distance']], index=[before_best])
        else:
            return pd.Series(
                [before.at[before_best, 'distance'], after.at[after_best, 'distance']],
                index=[before_best, after_best])

    def get_closest_features_to_list(var, features, target_type):
        tmp = [get_closest_features(var, fea, target_type) for fea in features]
        res = []
        for fea, ser in zip(features, tmp):
            tmpi = ser.to_frame()
            tmpi['feature'] = fea
            res.append(tmpi)
        res = pd.concat(res)
        res.rename(columns={0: 'closest'}, inplace=True)
        return res

    ct = 'Venous EC'

    genes = ['Slc6a2']
    closest_peaks = get_closest_features_to_list(adata.var, genes, 'Peaks')
    peaks = list(closest_peaks.index)
    features = genes + peaks
    avg = np.zeros((len(cell_types), len(features)))
    for i, cst in enumerate(cell_types):
        adatai = adata_nor_gr[adata_nor_gr.obs['Cell Type'] == cst, features]
        avgi = adatai.X.mean(axis=0)
        avg[i] = avgi
    avg = pd.DataFrame(avg, index=cell_types, columns=features)
    avg_plt = avg / avg.max(axis=0)
    # Rename gene by adding coordinates
    avg_plt.rename(columns={
        genes[0]: '{:} ({:}:{:}-{:})'.format(
            genes[0],
            *adata.var.loc['Slc6a2'][['chromosome', 'start', 'end']].values),
        }, inplace=True)
    fig, ax = plt.subplots(figsize=(2.5, 4))
    sns.heatmap(avg_plt, ax=ax, yticklabels=True, xticklabels=True,
                cmap='copper')
    fig.get_axes()[1].set_ylabel('Relative expr/access')
    fig.tight_layout()

    peaks = (da_lineage_dict[ct].query('feature_types == "Peaks"')
                                .nlargest(40, 'score')
                                .index)

    closest_genes = get_closest_features_to_list(adata.var, peaks, 'Gene Expression')

    cell_types = lineages['Endothelial']
    genes = list(closest_genes.index)
    features = list(peaks) + genes
    avg = np.zeros((len(cell_types), len(features)))
    for i, cst in enumerate(cell_types):
        adatai = adata_nor_gr[adata_nor_gr.obs['Cell Type'] == cst, features]
        avgi = adatai.X.mean(axis=0)
        avg[i] = avgi
    avg = pd.DataFrame(avg, index=cell_types, columns=features)

    # Reorder
    features_order = []
    for peak in peaks:
        features_order.append(peak)
        features_order.extend(
            list(closest_genes.loc[closest_genes['feature'] == peak].index))
    avg = avg.loc[:, features_order]

    # Use standard scale instead of Z scale because negative numbers are
    # meaningless, esp for ATAC
    #avg_plt = (avg - avg.mean(axis=0)) / avg.std(axis=0)
    avg_plt = avg / avg.max(axis=0)

    from scipy.stats import pearsonr
    corrs = []
    for gene, peak in closest_genes['feature'].items():
        r = pearsonr(avg[gene].values, avg[peak].values)[0]
        corrs.append(r)
    corrs = pd.Series(corrs, index=closest_genes.index, name='corr').to_frame()
    corrs['feature'] = closest_genes['feature']

    tmp = corrs.fillna(0).nlargest(8, 'corr')['feature']
    tmpi = set(list(tmp.index) + list(tmp.values))
    features_selected = [x for x in features_order if x in tmpi]
    closest_genes_selected = closest_genes.loc[tmp.index]
    avg_plt = avg_plt[features_selected]

    fig, ax = plt.subplots(figsize=(8, 4))
    sns.heatmap(avg_plt, ax=ax, yticklabels=True, xticklabels=True)
    #ax2 = ax.twiny()
    #ax2.set_xticks(ax.get_xticks())
    #xticks2 = []
    #for fea in features_order:
    #    if fea.startswith('chr'):
    #        xtick = '-'
    #    else:
    #        xtick = str(closest_genes.loc[fea, 'closest'])
    #    xticks2.append(xtick)
    #ax2.set_xticklabels(xticks2, rotation=90)
    for feature, grp in closest_genes_selected.groupby('feature'):
        i0 = features_selected.index(feature)
        i1 = i0 + len(grp) + 1
        rect = plt.Rectangle((i0, 0), i1 - i0, len(cell_types),
                             edgecolor='k', lw=2, facecolor='none',
                             )
        ax.add_patch(rect)
    fig.tight_layout()


import argparse
import json
import os
import sys

import socks
from telethon import TelegramClient

from tools.down_file import down_group
from tools.monit import StartMonit
from tools.tool import print_all_channel, Hook, print_group, initDb, md5
from tools.upload_file import upload_file

# 1.创建解释器
parser = argparse.ArgumentParser()
# 2.创建一个互斥参数组
mutex_group = parser.add_mutually_exclusive_group()
# 3.添加需要的参数
parser.add_argument('-c', '--config', default='config.json', help='配置文件')
parser.add_argument('-re', '--refresh', action='store_true', help='刷新缓存')
mutex_group.add_argument('-up', '--upload', action='store_true', help='上传文件')
mutex_group.add_argument('-down', '--download', action='store_true', help='下载文件')
mutex_group.add_argument('-print', action='store_true', help='打印消息')
mutex_group.add_argument('-m', '--monit', action='store_true', help='监控频道')
parser.add_argument('-id', help='频道ID')
parser.add_argument('-user', help='指定下载用户', default=None)
parser.add_argument('--range', default='>0', help='下载范围')
parser.add_argument('-path', help='上传路径')
parser.add_argument('-dau', default='N', choices=['y', 'Y', 'n', 'N'], help='上传完成删除原文件')
# 3.进行参数解析
args = parser.parse_args()
config_path = args.config
# 配置处理开始
with open(config_path, 'r', encoding='utf-8') as f:
    config = json.load(f)
api_id = config.get('api_id')
api_hash = config.get('api_hash')
phone = config.get('phone')
bot_token = config.get('bot_token')
if (phone is not None and bot_token is not None) or (phone is None and bot_token is None):
    print('请确认使用机器人登录还是电话号码登录')
    exit()
if phone:
    md5Token = md5(phone)
else:
    md5Token = md5(bot_token)
initDb(md5Token)
os.environ['save_path'] = save_path = config.get('save_path')
proxy_ip = config.get('proxy_ip')
proxy_port = config.get('proxy_port')
if proxy_port is not None:
    if proxy_ip is None:
        proxy_ip = '127.0.0.1'
    client = TelegramClient(md5Token, api_id, api_hash, proxy=(socks.SOCKS5, proxy_ip, proxy_port))
else:
    client = TelegramClient(md5Token, api_id, api_hash)

# 配置处理结束

async def upDate_dialogs():
    await client.get_dialogs()

# 展示登陆的信息
def show_my_inf(me):
    print("-----****************-----")
    print("Name:", me.username)
    print("ID:", me.id)
    print("-----login successful-----")

async def client_main():
    print("-client-main-")
    me = await client.get_me()
    show_my_inf(me)

if __name__ == '__main__':
    # 除了刷新缓存，都需要频道ID
    if not args.refresh and args.id is None:
        sys.exit(1)
    if args.upload and args.path is None:
        print('缺失上传路径')
        sys.exit(1)
    with client.start(phone=phone, bot_token=bot_token):
        client.loop.run_until_complete(client_main())
        client.loop.run_until_complete(Hook(client))
        if args.refresh:
            print_all_channel(client=client)
        elif args.download:
            if 't.me' in args.id:
                tmpList = args.id.split('/')
                channel_id = tmpList[-2]
                plus_func = '=' + tmpList[-1]
            else:
                channel_id = args.id
                plus_func = args.range
            for _id in channel_id.split('|'):
                client.loop.run_until_complete(down_group(client, _id, plus_func, args.user))
        elif args.upload:
            del_after_upload = True if args.dau.upper() == 'Y' else False
            client.loop.run_until_complete(upload_file(client, args.id, args.path, del_after_upload))
        elif args.print:
            client.loop.run_until_complete(print_group(client, args.id))
        elif args.monit:
            channel_ids = args.id.split(',')
            client.loop.run_until_complete(StartMonit(client, channel_ids))
            client.run_until_disconnected()


from botocore.exceptions import NoCredentialsError
from django.test import TestCase
from unittest.mock import (
    Mock,
    patch,
)
import json

from django.conf import settings

from common.common_utils.aws_sqs_utils import (
    create_sqs_event,
    get_sqs_messages,
)

class TestCreateSQSEvent(TestCase):
    @patch('common.common_utils.aws_sqs_utils.boto3.client')
    def test_create_sqs_event(self, mock_boto3_client):
        # Given: Set up the test data
        queue_url = 'https://sqs.us-east-1.amazonaws.com/123456789012/test-queue'
        event_data = {'key': 'value'}

        # And: Mock the boto3.client and sqs.send_message calls
        mock_sqs_client = Mock()
        mock_boto3_client.return_value = mock_sqs_client
        mock_send_message_response = {'MessageId': '1234567890'}
        mock_sqs_client.send_message.return_value = mock_send_message_response

        # When: Call the function
        response = create_sqs_event(queue_url, event_data)

        # Then: Assert that sqs.send_message is called with the correct data
        mock_boto3_client.assert_called_once_with(
            'sqs',
            region_name='ap-northeast-2',
            aws_access_key_id=settings.AWS_IAM_ACCESS_KEY,
            aws_secret_access_key=settings.AWS_IAM_SECRET_ACCESS_KEY,
        )
        mock_sqs_client.send_message.assert_called_once_with(
            QueueUrl=queue_url,
            MessageBody=json.dumps(event_data),
        )

        # And: Assert the response from the function
        self.assertEqual(response, mock_send_message_response)

class TestGetSQSMessages(TestCase):
    @patch('common.common_utils.aws_sqs_utils.boto3.client')
    def test_get_sqs_messages(self, mock_boto3_client):
        # Given: Set up the test data
        queue_url = 'https://sqs.us-east-1.amazonaws.com/123456789012/test-queue'
        message_body = '{"key": "value"}'

        # And: Mock the boto3.client and sqs.receive_message calls
        mock_sqs_client = Mock()
        mock_boto3_client.return_value = mock_sqs_client
        mock_receive_message_response = {
            'Messages': [
                {
                    'Body': message_body,
                    'ReceiptHandle': 'mock-receipt-handle'
                }
            ]
        }
        mock_sqs_client.receive_message.return_value = mock_receive_message_response

        # When: Call the function
        result = get_sqs_messages(queue_url)

        # Then: Assert that sqs.receive_message and sqs.delete_message are called with the correct data
        mock_boto3_client.assert_called_once_with(
            'sqs',
            region_name='ap-northeast-2',
            aws_access_key_id=settings.AWS_IAM_ACCESS_KEY,
            aws_secret_access_key=settings.AWS_IAM_SECRET_ACCESS_KEY,
        )
        mock_sqs_client.receive_message.assert_called_once_with(
            QueueUrl=queue_url,
            MaxNumberOfMessages=1
        )
        mock_sqs_client.delete_message.assert_called_once_with(
            QueueUrl=queue_url,
            ReceiptHandle='mock-receipt-handle'
        )

        # And: Assert the result
        expected_result = {'key': 'value'}
        self.assertEqual(result, expected_result)

    @patch('common.common_utils.aws_sqs_utils.boto3.client', side_effect=NoCredentialsError())
    def test_get_sqs_messages_with_exception(self, mock_boto3_client):
        # Given: Set up the test data
        queue_url = 'https://sqs.us-east-1.amazonaws.com/123456789012/test-queue'

        # Expected: Call the function and expect an exception
        with self.assertRaises(Exception):
            get_sqs_messages(queue_url)


# Imports
import tensorflow as tf
import os
import cv2
import numpy as np
import matplotlib.pyplot as plt

def load_data(directory):
    directories = [d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))]
    labels = []
    images = []
    print(directories)
    for d in directories:
        subdir = os.path.join(directory, d)
        print(subdir)
        filenames = [ os.path.join(subdir, file_name) for file_name in os.listdir(subdir) if file_name.endswith('.ppm')]
        for file_name in filenames:
            images.append(cv2.imread(file_name))
            labels.append(int(d))
    return images, labels

Root = '/data/Belgium_TSC/'
train_directory = os.path.join(Root, 'Training')
images, labels = load_data(train_directory)
images = np.array(images)
labels = np.array(labels)

images28 = [cv2.resize(image, (28, 28), interpolation = cv2.INTER_AREA) for image in images]
images28 = [cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) for image in images28]

model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28,1)),
  tf.keras.layers.MaxPooling2D((2, 2)),
  tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
  tf.keras.layers.MaxPooling2D((2, 2)),
  tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
  tf.keras.layers.Flatten(),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dense(62, activation='softmax')
])
model.summary()
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
images28 = np.array(images28)
images_extend = np.expand_dims(images28, -1)
model.fit(images_extend, np.array(labels), epochs=100)


import unittest
import trio
import os
import subprocess

from minio import Minio
from util.test_base import ApolloTest
from util import skvbc as kvbc
from util import eliot_logging as log
from util.object_store import ObjectStore, start_replica_cmd_prefix
from util.bft import KEY_FILE_PREFIX, with_trio, with_bft_network

def start_replica_cmd(builddir, replica_id, config):
    """
    There are two test s3 config files. The one used here hasn't got s3-prefix parameter
    set its value will be an empty string.
    The effect of this is that on each RO replica start we have got empty S3 storage.
    This is the default behaviour for most of the tests.
    """
    return start_replica_cmd_prefix(builddir, replica_id, config)

class SkvbcS3IntegrityCheckTest(ApolloTest):
    """
    Integrity Check has got two modes of operation:
        - validate all
        - validate key
    """

    @classmethod
    def setUpClass(cls):
        cls.object_store = ObjectStore()

    @classmethod
    def tearDownClass(cls):
        pass

    @with_trio
    @unittest.skip("unstable scenario")
    @with_bft_network(start_replica_cmd=start_replica_cmd_prefix, num_ro_replicas=1, selected_configs=lambda n, f, c: n == 7)
    async def test_integrity_check_validate_all(self, bft_network):
        """
        Start all replicas including ROR
        Generate a checkpoint
        start integrity
        check that integrity was verified
        This test is executed only in S3 mode.
        """
        bft_network.start_all_replicas()
        ro_replica_id = bft_network.config.n

        bft_network.start_replica(ro_replica_id)
        skvbc = kvbc.SimpleKVBCProtocol(bft_network)

        await skvbc.fill_and_wait_for_checkpoint(
            initial_nodes=bft_network.all_replicas(),
            num_of_checkpoints_to_add=1,
            verify_checkpoint_persistency=False,
            assert_state_transfer_not_started=False
        )

        await self._wait_for_st(bft_network, ro_replica_id)

        keys_config = f"{KEY_FILE_PREFIX}{ro_replica_id}"
        s3_config = "test_s3_config_prefix.txt"
        self._start_integrity_check(bft_network, keys_config, s3_config)

    @with_trio
    @unittest.skip("unstable scenario")
    @with_bft_network(start_replica_cmd=start_replica_cmd_prefix, num_ro_replicas=1, selected_configs=lambda n, f, c: n == 7)
    async def test_integrity_check_validate_all_with_missing_block(self, bft_network):
        """
        Start all replicas including ROR
        Generate a checkpoint
        start integrity
        delete block
        verify that integrity check failed
        This test is executed only in S3 mode.
        """
        bft_network.start_all_replicas()
        ro_replica_id = bft_network.config.n

        bft_network.start_replica(ro_replica_id)
        skvbc = kvbc.SimpleKVBCProtocol(bft_network)

        await skvbc.fill_and_wait_for_checkpoint(
            initial_nodes=bft_network.all_replicas(),
            num_of_checkpoints_to_add=1,
            verify_checkpoint_persistency=False,
            assert_state_transfer_not_started=False
        )

        await self._wait_for_st(bft_network, ro_replica_id)
        self._delete_block()
        keys_config = f"{KEY_FILE_PREFIX}{ro_replica_id}"
        s3_config = "test_s3_config_prefix.txt"
        try:
            self._start_integrity_check(bft_network, keys_config, s3_config)
            assert False
        except AssertionError as e:
            pass

    def _start_integrity_check(self, bft_network, keys_file, s3_config_file, key_to_validate=None):
        """
        Start integrity check
        """
        with log.start_action(action_type="start_integrity_check"):
            stdout_file = None
            stderr_file = None

            if os.environ.get('KEEP_APOLLO_LOGS', "").lower() in ["true", "on"]:
                test_name = os.environ.get('TEST_NAME')

                if not test_name:
                    now = datetime.now().strftime("%y-%m-%d_%H:%M:%S")
                    test_name = f"{now}_{bft_network.current_test}"

                test_dir = f"{bft_network.builddir}/tests/apollo/logs/{test_name}/{bft_network.current_test}/"
                test_log = f"{test_dir}stdout_integrity_check.log"
                log.log_message(message_type=f"test log is: {test_log}")
                os.makedirs(test_dir, exist_ok=True)

                stdout_file = open(test_log, 'w+')
                stderr_file = open(test_log, 'w+')

                stdout_file.write("############################################\n")
                stdout_file.flush()
                stderr_file.write("############################################\n")
                stderr_file.flush()

                s3_config_path = os.path.join(bft_network.builddir, s3_config_file)
                integrity_check_fds = (stdout_file, stderr_file)
                integrity_check_exe = os.path.join(bft_network.builddir, "kvbc", "tools", "object_store_utility", "object_store_utility")
                integrity_check_cmd = [integrity_check_exe,
                                       "--keys-file", keys_file,
                                       "--s3-config-file", s3_config_path]
                integrity_check_cmd.append("validate")
                if key_to_validate is not None:
                    integrity_check_cmd.append("--key")
                    integrity_check_cmd.append(key_to_validate)
                else:
                    integrity_check_cmd.append("--all")
                log.log_message(message_type="starting the subprocess")
                integrity_check_pid = subprocess.Popen(
                    integrity_check_cmd,
                    stdout=stdout_file,
                    stderr=stderr_file,
                    close_fds=True)
                try:
                    exit_code = integrity_check_pid.wait()
                    assert exit_code == 0
                except Exception as e:
                    assert False
                finally:
                    for fd in integrity_check_fds:
                        fd.close()

    def _delete_block(self):
        client = Minio("127.0.0.1:9000", access_key="concordbft", secret_key="concordbft", secure=False)
        buckets = client.list_buckets()
        bucketName = buckets[0].name
        try:
            client.remove_object(bucketName, 'concord/blocks/1')
        except ResponseError as err:
            assert False

    async def _wait_for_st(self, bft_network, ro_replica_id, seqnum_threshold=150):
        with trio.fail_after(seconds=70):
            # the ro replica should be able to survive these failures
            while True:
                with trio.move_on_after(seconds=.5):
                    try:
                        key = ['replica', 'Gauges', 'lastExecutedSeqNum']
                        lastExecutedSeqNum = await bft_network.metrics.get(ro_replica_id, *key)
                    except KeyError:
                        continue
                    else:
                        # success!
                        if lastExecutedSeqNum >= seqnum_threshold:
                            log.log_message(message_type="Replica" + str(ro_replica_id) + " : lastExecutedSeqNum:" + str(lastExecutedSeqNum))
                            break


import discord
from discord.ext import commands, tasks

bot = commands.Bot (command_prefix="!")

token = "토큰"

@bot.event
async def on_ready():
    print("Bot Connected")
    print(f"Bot name: {bot.user.name}")
    print(f"ID: {bot.user.id}")

@bot.command()
async def ping(ctx):
    await ctx.send(":ping_pong:  Pong!")

bot.run(token)


def combine(n: int, k: int):
    base = []
    def permu(step, base):
            if step == 0:
                return base
            else:
                check = []
                for i in base:
                    temp = i[:]
                    for j in range(i[-1]+1, n+1):
                        temp.append(j)
                        check.append(temp)
                        if j < n:
                            temp = i[:]
                return permu(step-1, check)
    for i in range(1, n-k+2):
        base.append([i])
    return permu(k-1, base)

def combine2(n: int, k: int) :
        ans = []
        def helper(combination, start, remaining_slots):
            if remaining_slots == 0:
                ans.append(combination)
            else:
                for num in range(start, n + 1):
                    helper(combination + [num], num + 1, remaining_slots - 1)
                    print(combination)
        helper([], 1, k)
        return ans

print(combine(5, 3))

def combine3(self, n, k):
        ret = []
        self.dfs(list(range(1, n+1)), k, [], ret)
        return ret

def dfs(self, nums, k, path, ret):
        if len(path) == k:
            ret.append(path)
            return
        for i in range(len(nums)):
            self.dfs(nums[i+1:], k, path+[nums[i]], ret)


import tkinter
import tkinter.messagebox

def displaymessage():
    tkinter.messagebox.showinfo("Python GUI", "SUBMITTED!!")

window1 = tkinter.Tk()  # define canvas
window1.geometry("300x300")
#add a button
button1 = tkinter.Button(window1, text='Start', bg='cyan', command=displaymessage)
button1.config(font=("Courier,50"))

button1.pack(side='right', expand=tkinter.YES)

window1.mainloop()


# %%
import json
import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from PIL import Image
import PIL
from glob import glob
import decord
from decord import VideoReader
import random
import logging

try:
    from lavis.datasets.datasets.base_dataset import BaseDataset
except ModuleNotFoundError:
    import sys
    sys.path.append(os.path.realpath(os.path.join(__file__, '../LAVIS')))
    from lavis.datasets.datasets.base_dataset import BaseDataset

from torch.utils.data.dataloader import default_collate
import pandas as pd
from collections import defaultdict
from typing import Literal, get_args, List

logging.getLogger().setLevel(logging.INFO)

decord.bridge.set_bridge('native')

### === ActionBenchDataset === ###
class ActionBenchDatasetBase(BaseDataset):
    def __init__(self,
        annotation,
        task = "video_text_matching",
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps=None,
        k=None,
        neg_sampling_same_clip=0
    ):
        """
            annotation: a list of annotation samples
        """
        self.task = task
        self.annotation = annotation
        print("annotation size:", len(self.annotation))

        self.vis_processor = vis_processor
        self.text_processor = text_processor

        self.frm_sampling_strategy = frm_sampling_strategy
        self.frame_height = frame_height
        self.frame_width = frame_width
        self.num_frm = num_frm
        self.fps = fps
        print("using fps:", self.fps)

        # if using negative sampling
        self.neg_sampling_same_clip = neg_sampling_same_clip
        if self.neg_sampling_same_clip > 0:
            print("Using negative sampling from same clip, number:", neg_sampling_same_clip)

    def _load_video_from_path_decord(self, video_path, fps, start_time=None, end_time = None):
        """
            returns frames as PIL images
        """
        frm_sampling_strategy = self.frm_sampling_strategy
        num_frm = self.num_frm
        task = self.task
        try:
            # load video
            if not self.frame_height or not self.frame_width:
                vr = VideoReader(video_path)
            else:
                vr = VideoReader(video_path, width=self.frame_width, height=self.frame_height)

            # get start and end frame idx
            vlen = len(vr)
            if start_time or end_time:
                assert fps > 0, 'must provide video fps if specifying start and end time.'
                start_idx = max(int(start_time * fps), 0)
                end_idx = min(int(end_time * fps), vlen-1)
                if start_idx == end_idx:
                    end_idx = min(start_idx + 1, vlen-1)
            else:
                start_idx, end_idx = 1, vlen-1

            # task

            if frm_sampling_strategy == 'uniform':
                frame_indices = np.linspace(
                    start_idx,
                    end_idx,
                    num = num_frm,
                    endpoint = True,
                    retstep = False,
                    dtype = int
                )
            else:
                raise NotImplementedError('Invalid sampling strategy {} '.format(frm_sampling_strategy))

            raw_sample_frms = vr.get_batch(frame_indices).asnumpy() # (num_frm, H, W, C)
            # raw_sample_frms = vr.get_batch(frame_indices) # tensor (num_frm, H, W, C)
            raw_sample_frms = [Image.fromarray(item, mode="RGB") for item in raw_sample_frms] # PIL Images

        except Exception as e:
            logging.info(e)
            return None
        # raw_sample_frms = raw_sample_frms.permute(0, 3, 1, 2) # torch tensor
        # raw_sample_frms = transforms.ToPILImage()(raw_sample_frms).convert('RGB')
        return raw_sample_frms

    def collater(self, samples, additional_fields: List[str] = [], text_fields: List[str] = []):
        if self.vis_processor is None:
            return samples
        if self.neg_sampling_same_clip == 0:
            return default_collate(samples)
        else:
            # additional work for batching the positive and negative samples
            ret = {
                "video_input":[],
                "video_input_reversed":[],
                "video_input_before":[],
                "video_input_after":[],
                "text_input":[],
                "action_antonym_text_input":[],
                "text_input_before":[],
                "text_input_after":[],
                **{f : [] for f in additional_fields}
            }
            for sample in samples:
                pos_instance = sample['pos_instance']
                neg_instances = sample['neg_instances']
                for att in ret.keys():
                    if att in pos_instance:
                        ret[att] += [pos_instance[att]] + [neg[att] for neg in neg_instances]

            empty_field_keys = []
            for att in ret.keys():
                if att not in text_fields:
                    if ret[att] != []:
                        ret[att] = torch.stack(ret[att])
                if ret[att] == []:
                    empty_field_keys.append(att)
            for remove_key in empty_field_keys:
                del ret[remove_key]
            return ret

    def __getitem__(self, index):
        if self.neg_sampling_same_clip == 0:
            # if not using any negative sampling strategy
            return self._get_single_item(index)
        else:
            # additional work if using the negative sampling strategy from the same clip
            ann_instance = self.annotation[index]
            indices_from_same_clip = list(self.clipuid_2_idx[ann_instance['clip_uid']])
            indices_from_same_clip.remove(index)
            assert len(indices_from_same_clip) == len(self.clipuid_2_idx[ann_instance['clip_uid']]) - 1
            if len(indices_from_same_clip) == 0:
                # the clip has only this one instance, then randomly sample another instance as negative
                sampled_neg_indices = random.sample(list(range(len(self.annotation))), self.neg_sampling_same_clip)
                print("Warning: no other instances from the same clip can be sampled, randomly sample negative instances:", sampled_neg_indices)
            else:
                while len(indices_from_same_clip) < self.neg_sampling_same_clip:
                    indices_from_same_clip += indices_from_same_clip
                sampled_neg_indices = random.sample(indices_from_same_clip, self.neg_sampling_same_clip)

            pos_instance = self._get_single_item(index) # Dict
            neg_instances = [self._get_single_item(idx) for idx in sampled_neg_indices] # List[Dict]
            return {
                "pos_instance":pos_instance,
                "neg_instances":neg_instances
            }

class ActionBenchDataset_Ego4D(ActionBenchDatasetBase):
    def __init__(self,
        vis_root,
        ann_path,
        state_change_filtering_json = None,
        state_change_filtering = False, # if true, use the filtered videos only during loading annotation
        task = "video_text_matching", # ["video_text_matching", "action_antonym", "object_shuffle", "reversed_video"]
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps=None,
        k=None,
        neg_sampling_same_clip=0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of EgoClip annotations
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/Ego4D/video_clips/clips_downsampled_5fps_downsized_224x224)
            ann_path (string): path to the egoclip style ann jsonl, each line is one instance (e.g., ActionBench/ego4d/egoclip_subset_action_antonyms_train_val_test_split/val.jsonl)
            state_change_filtering_json: path to the json file that contains indication of whether a video is state-change salient (e.g., ActionBench/ego4d/egoclip_subset_action_antonyms_train_val_test_split/state_change_heavy_instance_filtering_val.json)
            task (string): what's the evaluation task: choose from ["video_text_matching", "action_antonym", "object_shuffle", "reversed_video"]
                - "video_text_matching": for general vl tasks; returns a video_input and the corresponding text
                - "action_antonym": returns a video_input and two text input: one original annotation text, one modified annotation text with all verbs repleced with their antonym
                - "object_shuffle": returns a video_input and two text input: one original annotation text, one modified annotation text with action's objects repleced with another random object
                - "reversed_video": returns two video input: an original video_input and a reversed video_input with all frames reversed; and a text input of the original text annotation
                Note: all video_input returned are a sequence of sampled frames with size (num_frm, C, H, W)
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        # Ego4d nouns and verbs taxonomy size: see ActionBench/ego4d/*.csv for the index mapping
        self.noun_dim = 582  # num of nouns of ego4d taxonomy dictionary, not used for now, could be used for fine-grained negative sampling
        self.verb_dim = 118  # num of verbs of ego4d taxonomy dictionary, not used for now, could be used for fine-grained negative sampling

        self.vis_root = vis_root
        if isinstance(ann_path, list):
            ann_path = ann_path[0]

        self.clipuid_2_idx = defaultdict(list)

        if state_change_filtering_json is not None:
            self.state_change_filtering_dict = json.load(open(state_change_filtering_json))
        else:
            self.state_change_filtering_dict = None

        ann_full = []
        with open(ann_path) as f:
            for line in f:
                ann_instance = json.loads(line)
                clip_uid = ann_instance['clip_uid']
                clip_relative_start = ann_instance['clip_relative_start']
                clip_relative_end = ann_instance['clip_relative_end']
                # filter out missing video clips
                if os.path.exists(os.path.join(vis_root, f'{clip_uid}.mp4')):
                    # check if filter
                    if self.state_change_filtering_dict is not None and state_change_filtering:
                        key_ = clip_uid + "__" + str(clip_relative_start) + "__" + str(clip_relative_end)
                        if self.state_change_filtering_dict[key_]["if_keep"]:
                            self.clipuid_2_idx[clip_uid].append(len(ann_full))
                            ann_full.append(ann_instance)
                    else:
                        self.clipuid_2_idx[clip_uid].append(len(ann_full))
                        ann_full.append(ann_instance)

                    if k is not None and len(ann_full) >= k:
                        break

        super().__init__(
            ann_full,
            task = task, # ["video_text_matching", "action_antonym", "object_shuffle", "reversed_video"]
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps=fps,
            k=k,
            neg_sampling_same_clip=neg_sampling_same_clip
        )

        print("annotation size:", len(self.annotation))

    def _get_noun_verb_vec(self, sample):
        noun_vec = torch.zeros(self.noun_dim)
        verb_vec = torch.zeros(self.verb_dim)
        if 'tag_noun' in sample and 'tag_verb' in sample:
            noun_idx = eval(sample['tag_noun'])
            verb_idx = eval(sample['tag_verb'])
            for i in noun_idx:
                noun_vec[i] = 1
            for i in verb_idx:
                verb_vec[i] = 1
        else:
            noun_vec = None
            verb_vec = None
        return noun_vec, verb_vec # if output zero tensor, meaning not used for deciding loss mask

    def _get_single_item(self, index):
        """otuput format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann_instance = self.annotation[index]
        narration_text = ann_instance['clip_text']

        clip_uid = ann_instance['clip_uid']
        video_path = os.path.join(self.vis_root, f"{clip_uid}.mp4")
        start_time = ann_instance['clip_relative_start']
        end_time = ann_instance['clip_relative_end']
        if self.fps is None:
            fps = ann_instance['clip_fps']
        else:
            fps = self.fps # overiding the original fps in ann

        # get state_change_saliency_flag
        key_ = clip_uid + "__" + str(start_time) + "__" + str(end_time)
        if self.state_change_filtering_dict is not None:
            state_change_saliency_flag = self.state_change_filtering_dict[key_]["if_keep"]
        else:
            state_change_saliency_flag = True
        state_change_saliency_flag = torch.tensor(state_change_saliency_flag)

        # get video frames
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            fps,
            start_time=start_time,
            end_time=end_time
        )

        noun_vec, verb_vec = self._get_noun_verb_vec(ann_instance) # potentially useful for negative sampling (Egoclip annotation specific)

        if raw_sampled_frames is None:
            print(f"ERROR: cannot loading video clip: {ann_instance['clip_uid']}; try load another random instance...")
            return self.__getitem__(index+1)

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        if self.task == "action_antonym":
            narration_text_action_antonym = ann_instance['action_antonym_clip_text']
            if self.text_processor is not None:
                narration_text_action_antonym = self.text_processor(narration_text_action_antonym)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "action_antonym_text_input": narration_text_action_antonym,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        elif self.task == "object_shuffle":
            narration_text_object_shuffled = ann_instance['object_shuffled_clip_text']
            if self.text_processor is not None:
                narration_text_object_shuffled = self.text_processor(narration_text_object_shuffled)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "object_shuffled_text_input": narration_text_object_shuffled,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        elif self.task == "reversed_video":
            raw_sampled_frames_reverse = torch.flip(raw_sampled_frames, dims=[0])
            return {
                "video_input_reversed": raw_sampled_frames_reverse,
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        else: # "video_text_matching"
            # default output for general video-text matching tasks
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "noun_vec": noun_vec,
                "verb_vec": verb_vec,
                "state_change_saliency_flag": state_change_saliency_flag
            }

    ## overwrite collater here
    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=['noun_vec', 'verb_vec', "object_shuffled_text_input", "state_change_saliency_flag"],
            text_fields=["text_input","action_antonym_text_input","text_input_before","object_shuffled_text_input","text_input_after"]
        )

Split = Literal['train', 'val', 'test']
class ActionBenchDataset_SSv2(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: str,
        ann_root: str,
        split: Split,
        state_change_filtering = False, # if True, only keep the clips with state-change saliency
        use_templates_as_labels = False, # if True, use templates instead of the original annotations for text_input
        task = "video_text_matching", # ["video_text_matching", "action_antonym", "reversed_video"]
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm =  8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/SSv2/video_clips/clips_downsampled_5fps_downsized_224x224)
            ann_root (string): path to the annotations root directory (e.g., ActionBench/ssv2/shuffled_object_and_action_antonyms)
            use_templates_as_labels (bool): Whether to output the raw label template instead of the in-filled template. The test set only has raw templates as annotations.
            task (string): what's the evaluation task: choose from ["video_text_matching", "action_antonym", "reversed_video"]
                - "video_text_matching": for general vl tasks; returns a video_input and the corresponding text
                - "action_antonym": returns a video_input and two text input: one original annotation text, one modified annotation text with all verbs repleced with their antonym
                - "reversed_video": returns two video input: an original video_input and a reversed video_input with all frames reversed; and a text input of the original text annotation
                Note: all video_input returned are a sequence of sampled frames with size (num_frm, C, H, W)
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        if split not in get_args(Split):
            raise ValueError(f'split must be one of {get_args(Split)}')

        if split == 'test' and not use_templates_as_labels:
            raise ValueError('Test set only has raw templates as annotations')

        self.split = split

        self.vis_root = vis_root
        self.ann_root = ann_root
        self.use_templates_as_labels = use_templates_as_labels

        self.templates_to_ints = self._load_templates_to_ints()

        self.task = task # _load_annotations needs this to be set
        self.state_change_filtering = state_change_filtering # _load_annotations needs this to be set
        annotations = self._load_annotations(split)

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task, # ["video_text_matching", "action_antonym", "reversed_video"]
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_templates_to_ints(self):
        '''
        Loads mapping from templates to integers
        '''
        with open(os.path.join(self.ann_root, 'labels.json')) as f:
            j = json.load(f)

        return  {k : int(v) for k, v in j.items()}

    def _load_annotations(self, split: Split):
        '''
        Builds list of annotations each containing a video id, label, template, and placeholders
        from the annotation files.

        Args:
            split (Split): The stage to load the annotations for.

        Returns: A list of annotations with the following structure:
            [
                {
                    'clip_uid': int,
                    'label' : Optional[str],
                    'template' : str,
                    'placeholders' : Optional[List[str]]
                }
            ]
        '''
        # load train_antonym if exists; for look up action antonym for the original training set
        # for training instances that doesn't have an antonym, we will randomly pick another instance as action antonym text
        id_to_train_antonym_anns = {}
        train_antonym_path = os.path.join(self.ann_root, f'train_antonym.json')
        if os.path.exists(train_antonym_path):
            train_antonym_anns = json.load(open(train_antonym_path))
            for item in train_antonym_anns:
                id_to_train_antonym_anns[item['id']] = item

        if split in ['train', 'val']:
            # load state_change_filtering_json if any
            state_change_filtering_json_path = os.path.join(self.ann_root, f"state_change_heavy_instance_filtering_{split}.json")
            if os.path.exists(state_change_filtering_json_path):
                self.state_change_filtering_dict = json.load(open(state_change_filtering_json_path))
            else:
                self.state_change_filtering_dict = None

            # load main ann
            split = 'validation' if split == 'val' else 'train' # File is called "validation.json"

            with open(os.path.join(self.ann_root, f'{split}.json')) as f:
                j = json.load(f)

            annots = []
            for d in j:
                annot = {
                    'clip_uid': int(d['id']),
                    'label': d['label'].capitalize(),
                    'template': d['template'].replace('[something]', 'something'),
                    'placeholders': d['placeholders']
                }

                if self.state_change_filtering_dict is not None and self.state_change_filtering:
                    if not self.state_change_filtering_dict[str(d['id'])]['if_keep']:
                        continue

                if self.task == 'action_antonym':
                    if split == 'validation':
                        # assert that validation instances contains antonym items
                        annot.update({
                            'label_action_antonym_clip_text': d['label_action_antonym_clip_text'],
                            'template_action_antonym_clip_text': d['template_action_antonym_clip_text']
                        })
                    else: # split == 'train'
                        id_ = int(d['id'])
                        if id_ in id_to_train_antonym_anns:
                            # if current id has an action antonym
                            train_antonym_ann = id_to_train_antonym_anns[id_]
                            annot.update({
                                'label_action_antonym_clip_text': train_antonym_ann['label_action_antonym_clip_text'],
                                'template_action_antonym_clip_text': train_antonym_ann['template_action_antonym_clip_text']
                            })
                        else:
                            # if not, randomly pick an item in train_antonym, and using the original text as antonym
                            random_id = random.choice(list(id_to_train_antonym_anns.keys()))
                            train_antonym_ann = id_to_train_antonym_anns[random_id]
                            annot.update({
                                'label_action_antonym_clip_text': train_antonym_ann['label'],
                                'template_action_antonym_clip_text': train_antonym_ann['template']
                            })

                if self.task == 'object_shuffle':
                    annot.update({
                        'label_object_shuffled_clip_text': d['label_object_shuffled_clip_text']
                    })

                annots.append(annot)
        else:
            # columns = ['id', 'template'] if self.task not in ['action_antonym','object_shuffle','reversed_video'] else None # Infer from csv if action_antonym
            # df = pd.read_csv(os.path.join(self.ann_root, 'test-answers.csv'), names=columns, sep=';')

            df = pd.read_csv(os.path.join(self.ann_root, 'test-answers.csv'), header=0, sep=';')

            annots = []
            for r in df.itertuples():
                annot = {
                    'clip_uid': int(r.id),
                    'label': None,
                    'template': r.template,
                    'placeholders': None
                }

                if self.task == 'action_antonym':
                    annot.update({
                        'label_action_antonym_clip_text': None,
                        'template_action_antonym_clip_text': r.template_action_antonym_clip_text
                    })

                annots.append(annot)

        return annots

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['template'] if self.use_templates_as_labels else ann['label']

        id = ann['clip_uid']

        # get state_change_saliency_flag
        if self.state_change_filtering_dict is not None:
            state_change_saliency_flag = self.state_change_filtering_dict[str(id)]['if_keep']
        else:
            state_change_saliency_flag = True
        state_change_saliency_flag = torch.tensor(state_change_saliency_flag)

        # get video frames
        video_path = os.path.join(self.vis_root, f"{id}.mp4")
        start_time = None # ann_instance['clip_relative_start']
        end_time = None # ann_instance['clip_relative_end']

        assert os.path.exists(video_path)
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        if self.task == "action_antonym":
            narration_text_action_antonym = ann[f'{"template" if self.use_templates_as_labels else "label"}_action_antonym_clip_text']
            if self.text_processor is not None:
                narration_text_action_antonym = self.text_processor(narration_text_action_antonym)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "action_antonym_text_input": narration_text_action_antonym,
                "state_change_saliency_flag": state_change_saliency_flag
            }

        elif self.task == "object_shuffle":
            assert not self.use_templates_as_labels
            narration_text_object_shuffled = ann['label_object_shuffled_clip_text']
            if self.text_processor is not None:
                narration_text_object_shuffled = self.text_processor(narration_text_object_shuffled)
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "object_shuffled_text_input": narration_text_object_shuffled,
                "state_change_saliency_flag": state_change_saliency_flag
            }

        elif self.task == "reversed_video":
            raw_sampled_frames_reverse = torch.flip(raw_sampled_frames, dims=[0])
            return {
                "video_input_reversed": raw_sampled_frames_reverse,
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "state_change_saliency_flag": state_change_saliency_flag
            }
        else:
            # default output for general video-text matching tasks
            return {
                "video_input": raw_sampled_frames,
                "text_input": narration_text,
                "state_change_saliency_flag": state_change_saliency_flag
            }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["object_shuffled_text_input", "state_change_saliency_flag"],
            text_fields=["text_input","action_antonym_text_input","text_input_before","object_shuffled_text_input","text_input_after"]
        )
        # return super().collater(
        #     samples,
        #     text_fields=["text_input","action_antonym_text_input","text_input_before","text_input_after"]
        # )

### === Downstream Datasets === ###
class DownstreamTask_Retrieval_SSv2(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: str,
        ann_root: str,
        split: Literal['train', 'val'],
        task = Literal['ssv2_label','ssv2_template'],
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/SSv2/video_clips/clips_downsampled_5fps_downsized_224x224)
            ann_root (string): path to the annotations root directory (e.g., datasets/SSv2/ssv2_label_ssv2_template)
            use_templates_as_labels (bool): Whether to output the raw label template instead of the in-filled template. The test set only has raw templates as annotations.
            task (string): ssv2_label or ssv2_template as described in https://arxiv.org/abs/2206.03428
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = split

        self.vis_root = vis_root
        self.ann_root = ann_root

        # self.templates_to_ints = self._load_templates_to_ints()

        self.task = task # _load_annotations needs this to be set

        (
            annotations,
            self.videos,
            self.texts,
            self.v2t_targets,
            self.t2v_targets
        ) = self._load_annotations(split)

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_annotations(self, split):
        if split == "train" and self.task == "ssv2_label":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_label_train.json")
        elif split == "val" and self.task == "ssv2_label":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_label_val_small.json")
        elif split == "train" and self.task == "ssv2_template":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_template_train.json")
        elif split == "val" and self.task == "ssv2_template":
            ann_path = os.path.join(self.ann_root,"ssv2_ret_template_val_small.json")

        j = json.load(open(ann_path))

        annots = []
        videos = []
        texts = []
        v2t_targets = []
        t2v_targets = []
        for item in j:
            video_names = item['video']
            if isinstance(video_names, str):
                video_names = [video_names]
            caption = item['caption'].replace('[something]', 'something')
            gt_video_idx_this_caption = []
            for video_name in video_names:
                video_id = video_name[:-5]
                if not self._check_if_video_exists(video_id):
                    logging.info("!!!! video {} does not exist, skip...".format(video_id))
                    continue
                # if video exists
                gt_video_idx_this_caption.append(len(videos)) # add current video idx
                videos.append(video_id)
                v2t_targets.append([len(texts)]) # one video only have one gt text
                annot = {
                    'clip_uid': int(video_id),
                    'caption': caption
                }
                annots.append(annot)
            texts.append(caption)
            t2v_targets.append(gt_video_idx_this_caption)

        assert len(videos) == len(v2t_targets)
        assert len(texts) == len(t2v_targets)
        return annots, videos, texts, v2t_targets, t2v_targets

    def _check_if_video_exists(self, video_id):
        return os.path.exists(os.path.join(self.vis_root, f"{video_id}.mp4"))

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['caption']

        id = ann['clip_uid']
        video_path = os.path.join(self.vis_root, f"{id}.mp4")
        start_time = None # ann_instance['clip_relative_start']
        end_time = None # ann_instance['clip_relative_end']

        assert os.path.exists(video_path)
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": narration_text,
            "idx": int(index),
        }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["idx"],
            text_fields=["text_input"]
        )

class DownstreamTask_MomentsInTime(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: str,
        ann_root: str,
        split: Literal['train', 'val'],
        task = Literal['video_action_retrieval','video_action_retrieval_2k'],
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (string): Root directory of video clips (e.g., datasets/Moments_In_Time/videos)
            ann_root (string): path to the annotations root directory (e.g., datasets/Moments_In_Time/ann)
            task (string): video to action label
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = "training" if split == "train" else "validation"
        self.vis_root = os.path.join(vis_root, self.split)
        self.ann_root = ann_root

        self.task = task # only one task for now

        (
            annotations,
            self.videos,
            self.texts,
            self.v2t_targets,
            self.t2v_targets
        ) = self._load_annotations()

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_annotations(self):
        if self.split == "training":
            ann_path = os.path.join(self.ann_root, "trainingSet.csv")
        else:
            if self.task == "video_action_retrieval_2k":
                ann_path = os.path.join(self.ann_root, "validationSet_2k.csv")
            else:
                ann_path = os.path.join(self.ann_root, "validationSet.csv")

        # get label texts & index mapping
        label_txt = os.path.join(self.ann_root, "moments_categories.txt")
        self.text_2_index = {}
        texts = []
        with open(label_txt, 'r') as f:
            for line in f:
                l, idx = line.split(',')
                l = l.strip().replace('+',' ')
                texts.append(l)
                self.text_2_index[l] = int(idx.strip())
        self.index_2_text = {value:key for key,value in self.text_2_index.items()}

        annotations = []
        videos = []
        v2t_targets = []
        t2v_targets = [[] for i in range(len(texts))]

        with open(ann_path, 'r') as f:
            for line in f:
                video_name, label, res_1, res_2 = line.split(',')
                vid_path = os.path.join(self.vis_root, video_name)
                if os.path.exists(vid_path):
                    text = label.strip().replace('+',' ')
                    annotations.append(
                        {
                            'clip_uid': video_name,
                            'caption': text
                        }
                    )
                    assert text in self.text_2_index
                    t_idx = self.text_2_index[text]
                    v2t_targets.append([t_idx])
                    t2v_targets[t_idx].append(len(videos))
                    videos.append(video_name)

        assert len(videos) == len(v2t_targets)
        assert len(texts) == len(t2v_targets)
        return annotations, videos, texts, v2t_targets, t2v_targets

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['caption']
        label_index = self.text_2_index[narration_text]

        video_path = os.path.join(self.vis_root, ann['clip_uid'])
        start_time = None # ann_instance['clip_relative_start']
        end_time = None # ann_instance['clip_relative_end']

        assert os.path.exists(video_path)
        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": narration_text,
            "label_index": label_index,
            "idx": int(index),
        }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["idx","label_index"],
            text_fields=["text_input"]
        )

class DownstreamTask_Temporal(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: dict,
        ann_root: str,
        split: Literal['val'],
        task = Literal['v1.0','v1.0_2.4k'],
        subset = Literal['kinetics','ssv2'],
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (dict): for example {
                'kinetics':'datasets/Temporal/video_clips/kinetics400/clips_downsampled_5fps_downsized_224x224',
                'ssv2': 'datasets/SSv2/video_clips/clips_downsampled_5fps_downsized_224x224'
            }
            ann_root (string): path to the annotations root directory (e.g. datasets/Temporal/ann)
            task (string): provided two versions of val set https://arxiv.org/abs/2301.02074
            subset: a list specifing the sub-datasets to use ['kinetic', 'ssv2'] or ['kinetic'] or ['ssv2']
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = split
        self.vis_root = vis_root[subset]
        self.ann_root = ann_root

        self.subset = subset # _load_annotations needs this to be set
        self.task = task # _load_annotations needs this to be set

        (
            annotations,
            self.videos,
            self.texts,
            self.v2t_targets,
            self.t2v_targets
        ) = self._load_annotations(split)

        self.clipuid_2_idx = {d['clip_uid'] : [i] for i, d in enumerate(annotations)} # Each clip has only one annotation

        self.neg_sampling_same_clip = 0 # set neg_sampling_same_clip to 0 for SSv2 since there is only one annotation per video

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = 0
        )

    def _load_annotations(self, split):
        if split == "val":
            if self.task == "v1.0":
                ann_path = os.path.join(self.ann_root, "val-v1.0.csv")
            elif self.task == "v1.0_2.4k":
                ann_path = os.path.join(self.ann_root, "val-v1.0-2.4k.csv")

            annots = []
            videos = []
            texts = []

            df = pd.read_csv(ann_path, header=0)

            text_2_video_idx = defaultdict(list)
            for index, row in df.iterrows():
                dataset_name = row['dataset'].lower()
                if dataset_name != self.subset:
                    continue
                if self.subset == "ssv2":
                    video_id = row['video_id'].strip()
                    start_time = None
                    end_time = None
                else:
                    video_id = row['video_id'].split("/")[1].strip()[:-14]
                    time_string = row['video_id'].split("/")[1].strip()[-13:]
                    start_time = float(time_string.split('_')[0])
                    end_time = float(time_string.split('_')[1])

                if not self._check_if_video_exists(video_id):
                    logging.info("!!!! video {} does not exist, skip...".format(video_id))
                    continue
                text = row['text'].replace('[something]', 'something')
                text_2_video_idx[text].append(len(videos))
                videos.append(video_id)
                annots.append(
                    {
                        'clip_uid': video_id,
                        'caption': text,
                        'start_time': start_time,
                        'end_time': end_time,
                    }
                )

            v2t_targets = [[] for i in range(len(videos))]
            t2v_targets = []
            for text,video_indices in text_2_video_idx.items():
                for v in video_indices:
                    v2t_targets[v] = [len(texts)]
                t2v_targets.append(video_indices)
                texts.append(text)

            assert len(videos) == len(v2t_targets)
            assert len(texts) == len(t2v_targets)
            return annots, videos, texts, v2t_targets, t2v_targets
        else:
            raise NotImplementedError("Temporal dataset for training not implemented")

    def _check_if_video_exists(self, video_id):
        return os.path.exists(os.path.join(self.vis_root, f"{video_id}.mp4"))

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        narration_text = ann['caption']
        start_time = ann['start_time']
        end_time = ann['end_time']

        id = ann['clip_uid']
        video_path = os.path.join(self.vis_root, f"{id}.mp4")
        assert os.path.exists(video_path)

        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            narration_text = self.text_processor(narration_text)

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": narration_text,
            "idx": int(index),
        }

    def collater(self, samples):
        return super().collater(
            samples,
            additional_fields=["idx"],
            text_fields=["text_input"]
        )

class DownstreamTask_QA_NextQA(ActionBenchDatasetBase):
    def __init__(self,
        vis_root: dict,
        ann_root: str,
        split: Literal['train','val','test'],
        task = "mc", # multiple choice QA
        vis_processor = None,
        text_processor = None,
        frm_sampling_strategy = "uniform",
        num_frm = 8,
        frame_height = None,
        frame_width = None,
        fps = None,
        k = None,
        neg_sampling_same_clip = 0
    ):
        """
        dataset for loading video frames from video clips and preprocessed subset of SSv2 for downstream retrieval tasks
        params:
            vis_root (dict): path to videos (datasets/NextQA/video_clips/NExTVideo_downsampled_5fps_downsized_224x224)
            ann_root (string): path to the annotations root directory (datasets/NextQA/ann/nextqa)
            split: train, val, test
            vis_processor (Object class in processor.py): visual processor: add data augmentation and converting raw_sampled_frames (list of PIL images) returned from `self._load_video_from_path_decord` -> tensor
            text_processor (Object class in processor.py): textual processor: some basic processing of the input text (text string -> text string)
            frm_sampling_strategy (string): how do we sample frames from the clip: ['uniform']
            k (int): if not none, take a subset of k instances
        """
        self.split = split
        self.vis_root = vis_root
        self.ann_root = ann_root

        self.map_vid_vidorID = json.load(open(os.path.join(self.ann_root, "map_vid_vidorID.json")))
        annotations, clipuid_2_idx = self._load_annotations(split)

        self.clipuid_2_idx = clipuid_2_idx # each clip_uid can have multiple annotations

        super().__init__(
            annotations,
            task = task,
            vis_processor = vis_processor,
            text_processor = text_processor,
            frm_sampling_strategy = frm_sampling_strategy,
            num_frm = num_frm,
            frame_height = frame_height,
            frame_width = frame_width,
            fps = fps,
            k = k,
            neg_sampling_same_clip = neg_sampling_same_clip
        )

    def _load_annotations(self, split):
        csv = os.path.join(self.ann_root, f"{split}.csv")
        df = pd.read_csv(csv, header=0)

        annots = []
        clipuid_2_idx = defaultdict(list)
        for index, row in df.iterrows():
            vid = str(row['video'])
            if self._check_if_video_exists(vid):
                clipuid_2_idx[vid].append(len(annots))
                annots.append(
                    {
                        'clip_uid': vid,
                        'question': row['question'],
                        'answer': int(row['answer']),
                        'choices': [row[a] for a in ['a0','a1','a2','a3','a4']],
                        'action_antonym_choices': eval(row['action_antonym_choices']),
                        'qid': int(row['qid']),
                        'type': row['type'],
                        'start_time': None,
                        'end_time': None,
                    }
                )
        return annots, clipuid_2_idx

    def _check_if_video_exists(self, video_id):
        return os.path.exists(os.path.join(self.vis_root, self.map_vid_vidorID[video_id] + ".mp4"))

    def _get_single_item(self, index):
        """output format:
            - video: tensor (num_frm, C, H, W)
            - text: string
        """
        ann = self.annotation[index]
        concat_qa_text = [ann['question']+"? "+ ans for ans in ann['choices']]
        action_antonym_concat_qa_text = [ann['question']+"? "+ ans for ans in ann['action_antonym_choices']]
        answer = ann['answer']
        start_time = ann['start_time']
        end_time = ann['end_time']

        id = ann['clip_uid']
        video_path = os.path.join(self.vis_root, self.map_vid_vidorID[id] + ".mp4")
        assert os.path.exists(video_path)

        raw_sampled_frames = self._load_video_from_path_decord(
            video_path,
            self.fps,
            start_time=start_time,
            end_time=end_time
        )

        if raw_sampled_frames is None:
            print(f"ERROR: failed to load video clip with uid: {id}; trying to load another random instance...")
            return self.__getitem__(index+1)
        else:
            logging.debug(f'Successfully loaded clip with uid: {id}')

        # preprocess frames and text
        if self.vis_processor is not None:
            try:
                raw_sampled_frames = self.vis_processor(raw_sampled_frames) # (num_frm, C, H, W)
            except:
                raw_sampled_frames = [self.vis_processor(f) for f in raw_sampled_frames] # (num_frm, C, H, W)
                raw_sampled_frames = torch.stack(raw_sampled_frames, dim=0)

        if self.text_processor is not None:
            concat_qa_text = [self.text_processor(t) for t in concat_qa_text]
        if self.text_processor is not None:
            action_antonym_concat_qa_text = [self.text_processor(t) for t in action_antonym_concat_qa_text]

        # task specific output
        return {
            "video_input": raw_sampled_frames,
            "text_input": concat_qa_text,
            "action_antonym_text_input": action_antonym_concat_qa_text,
            "answer": int(answer),
            "idx": int(index),
        }

    def collater(self, samples, text_fields = ['text_input', 'action_antonym_text_input']):
        if self.vis_processor is None:
            return samples
        ret = {
            "video_input":[],
            "text_input":[],
            "action_antonym_text_input":[],
            "answer": [],
            "idx": []
        }
        if self.neg_sampling_same_clip == 0:
            for sample in samples:
                for att in ret.keys():
                    ret[att].append(sample[att])
        else:
            # additional work for batching the positive and negative samples
            for sample in samples:
                pos_instance = sample['pos_instance']
                neg_instances = sample['neg_instances']
                for att in ret.keys():
                    if att in pos_instance:
                        ret[att] += [pos_instance[att]] + [neg[att] for neg in neg_instances]
        empty_field_keys = []
        for att in ret.keys():
            if att not in text_fields:
                if ret[att] != []:
                    if not isinstance(ret[att][0], torch.Tensor):
                        ret[att] = [torch.tensor(item) for item in ret[att]]
                    ret[att] = torch.stack(ret[att])
            if ret[att] == []:
                empty_field_keys.append(att)
        for remove_key in empty_field_keys:
            del ret[remove_key]
        return ret