import traceback
import math
import os.path
import pandas as pd
import numpy as np
import datetime
import plotly.graph_objects as go
from statsmodels.tsa.stattools import adfuller
from statsmodels.tsa.ar_model import AutoReg
from sklearn.metrics import mean_squared_error
from pandas_profiling import ProfileReport
from financepy.utils import from_datetime
from financepy.products.bonds import *

root_path = 'd:/ProjectRicequant/fxincome/'
chinabond_ttm_labels = ['6m', '1y', '2y', '3y', '4y', '5y', '6y', '7y', '8y', '9y', '10y']
predicted_ttm_labels = ['pred_6m', 'pred_1y', 'pred_2y', 'pred_3y', 'pred_4y', 'pred_5y', 'pred_6y', 'pred_7y',
                        'pred_8y', 'pred_9y', 'pred_10y']
maturities = [0.5, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

def to_bond_freq(freq):
    if freq == 1:
        return FrequencyTypes.ANNUAL
    elif freq == 2:
        return FrequencyTypes.SEMI_ANNUAL
    elif freq == 4:
        return FrequencyTypes.QUARTERLY
    else:
        raise NotImplementedError("Unknown frequency")

def plot_bond_ylds(yld_date: datetime, bonds: list, ylds: list):
    plt.figure(figsize=(12, 6))
    title = yld_date.strftime('%Y%m%d')
    plt.title(f'CDB YTM @{title}')
    yld_date = from_datetime(yld_date.date())
    x = []
    for bond in bonds:
        years_to_maturity = (bond._maturity_date - yld_date) / 365
        x.append(years_to_maturity)
    bond_ylds_scaled = [n * 100 for n in ylds]
    plt.plot(x, bond_ylds_scaled, 'o')

    plt.xlabel('Time to Maturity (years)')
    plt.ylabel('Yield To Maturity (%)')
    plt.ylim((min(bond_ylds_scaled) - 0.3, max(bond_ylds_scaled) * 1.1))
    plt.grid(True)
    plt.savefig(f'./results/EXCEPTION_ytm_{title}.jpg', dpi=600)
    plt.close()

def plot_ns_params(df: pd.DataFrame):
    df = df.sort_values(by='date')
    first_date = df.date.iloc[0]
    last_date = df.date.iloc[len(df) - 1]
    title = f"{first_date.strftime('%Y%m%d')}-{last_date.strftime('%Y%m%d')}"
    df = df.rename(columns={'sum_beta1_beta2': 'beta1+beta2'})
    df.plot(kind='line', x='date', y=['beta1', 'beta2', 'beta1+beta2'], figsize=(15, 8))
    plt.grid(True)
    plt.legend(loc='upper left', bbox_to_anchor=(1, 1))
    plt.savefig(f'./results/ns_params_{title}.jpg')
    plt.close()

def preprocess():
    source = 'cdb_valuation.csv'
    result = 'cdb_valuation_processed.csv'
    df = pd.read_csv(os.path.join(root_path, source), parse_dates=['date', 'issue_date', 'maturity_date'],
                     encoding='gb2312')
    df = df[df.coupon_type == '固定利率']
    # df = df[(df.coupon_info == '附息') | (df.coupon_info == '到期一次还本付息')]
    #  Wind的付息次数不正确，到期一次还本付息的付息次数显示为0，应修改为1
    df.loc[df.coupon_info == '到期一次还本付息', 'coupon_freq'] = 1
    df_2019_2022 = df[(df.outstanding_balance >= 800) & (df.date.dt.year >= 2019)]
    df_2017_2018 = df[(df.outstanding_balance >= 500) & (df.date.dt.year >= 2017) & (df.date.dt.year <= 2018)]
    df_2015_2016 = df[(df.outstanding_balance >= 400) & (df.date.dt.year >= 2015) & (df.date.dt.year <= 2016)]
    df_2014 = df[(df.outstanding_balance >= 350) & (df.date.dt.year == 2014)]
    df_2013 = df[(df.outstanding_balance >= 300) & (df.date.dt.year == 2013)]
    df_2010_2012 = df[(df.outstanding_balance >= 250) & (df.date.dt.year >= 2010) & (df.date.dt.year <= 2012)]

    df = pd.concat([df_2010_2012, df_2013, df_2014, df_2015_2016, df_2017_2018, df_2019_2022])
    df = df[df.bond_code.str.len() == 9]
    df = df[(df.maturity_date - df.date).dt.days / 365 <= 10]  # only bonds with maturity <= 10 years
    df.to_csv(os.path.join(root_path, result), index=False, encoding='gb2312')

def check_csv():
    # source = 'cdb_valuation_processed.csv'
    # bond_df = pd.read_csv(root_path + source, parse_dates=['date', 'issue_date', 'maturity_date'],
    #                       encoding='gb2312')
    # bond_dates = set(bond_df.date.dt.date.to_list())
    # ytm_df = pd.read_csv(root_path + 'cdb_ytm.csv', parse_dates=['date'], encoding='gb2312')
    # ytm_dates = set(ytm_df.date.dt.date.to_list())
    # difference = list(ytm_dates - bond_dates)
    # difference.sort()
    # print(difference)
    # print(len(difference))
    # df = pd.read_csv('./results/ns_params_20100104_20220729.csv', parse_dates=['date'])
    # report = ProfileReport(df)
    # report.to_file('./results/ns_params_report.html')
    df = pd.DataFrame(
        {"value": np.random.randn(36)},
        index=pd.date_range("2011-01-01", freq="M", periods=36),
    )
    table = pd.pivot_table(
        df, index=df.index.month, columns=df.index.year, values="value", aggfunc="sum"
    )
    print(df)
    print(table)

def gen_curves(fit_type: str, start: datetime.datetime, end: datetime.datetime, plot: bool) -> pd.DataFrame:
    source = 'cdb_valuation_processed.csv'
    df_src = pd.read_csv(os.path.join(root_path, source), parse_dates=['date', 'issue_date', 'maturity_date'],
                         encoding='gb2312')
    if fit_type == 'ns':
        ns_stats = pd.DataFrame(columns=['date', 'beta1', 'beta2', 'sum_beta1_beta2', 'beta3', 'tau'])
    elif fit_type == 'nss':
        nss_stats = pd.DataFrame(
            columns=['date', 'beta1', 'beta2', 'sum_beta1_beta2', 'beta3', 'beta4', 'tau1', 'tau2'])
    else:
        raise NotImplementedError("Unknown curve fit type")
    ytm_date = start
    while ytm_date <= end:
        df = df_src[df_src.date == ytm_date]
        if len(df) == 0:  # Today is not a buisness day.
            ytm_date += datetime.timedelta(days=1)
            continue
        bonds = []
        ytms = []
        for bond in df.itertuples():
            issue_date = from_datetime(bond.issue_date)
            maturity_date = from_datetime(bond.maturity_date)
            coupon = bond.coupon / 100
            freq = to_bond_freq(bond.coupon_freq)
            accrual_type = DayCountTypes.ACT_ACT_ISDA
            ytm = bond.ytm / 100
            b = Bond(issue_date, maturity_date, coupon, freq, accrual_type)
            bonds.append(b)
            ytms.append(ytm)

        try:
            if fit_type == 'ns':
                ns_bounds = [(-1, -1, -1, 1.3), (1, 1, 1, 3)]
                ns_stats = ns_fit(bonds, ns_stats, ytm_date, ytms, ns_bounds, plot=plot)
            elif fit_type == 'nss':
                nss_bounds = [(0, -1, -1, -1, 1.3, 1.3), (1, 1, 1, 1, 3, 3)]
                nss_stats = nss_fit(bonds, nss_stats, ytm_date, ytms, nss_bounds, plot=plot)
            # curve_function = CurveFitPolynomial(5)
            # fitted_curve = BondYieldCurve(eval_date, bonds, ytms, curve_function)
            # fitted_curve.plot('CDB Yield Curve by Quintic Polynomial')
            # curve_function = CurveFitBSpline()
            # fitted_curve = BondYieldCurve(eval_date, bonds, ytms, curve_function)
            # fitted_curve.plot('CDB Yield Curve by B Spline')
            # curve_function = CurveFitBSpline(4)
            # fitted_curve = BondYieldCurve(eval_date, bonds, ytms, curve_function)
            # fitted_curve.plot('CDB Yield Curve by B Spline power of 4')

        except Exception as e:
            print(f'Exception:{str(e)} {ytm_date}')
            traceback.print_exc()
            plot_bond_ylds(ytm_date, bonds, ytms)

        ytm_date += datetime.timedelta(days=1)
    start = start.strftime('%Y%m%d')
    end = end.strftime('%Y%m%d')
    if fit_type == 'ns':
        ns_stats.to_csv(os.path.join('./results/', f'ns_params_{start}_{end}.csv'), index=False, encoding='utf-8')
        return ns_stats
    elif fit_type == 'nss':
        nss_stats.to_csv(os.path.join('./results/', f'nss_params_{start}_{end}.csv'), index=False, encoding='utf-8')
        return nss_stats
    else:
        raise NotImplementedError("Unknown curve fit type")

def ns_fit(bonds, ns_stats, ytm_date, ytms, bounds, plot=False):
    curve_function = CurveFitNelsonSiegel(bounds=bounds)
    fitted_curve = BondYieldCurve(from_datetime(ytm_date), bonds, ytms, curve_function)
    beta1 = curve_function._beta1
    beta2 = curve_function._beta2
    beta3 = curve_function._beta3
    tau = curve_function._tau
    ns_stats.loc[len(ns_stats)] = [ytm_date, beta1, beta2, beta1 + beta2, beta3, tau]
    if plot:
        today = ytm_date.strftime('%Y%m%d')
        fitted_curve.plot(f'CDB Yield Curve by Nelson Siegel @{today}')
        address = f'./results/images/ns_cdb_ytm_{today}.jpg'
        plt.savefig(address, dpi=600)
        plt.close()
    return ns_stats

def nss_fit(bonds, nss_stats, ytm_date, ytms, bounds, plot=False):
    bounds = bounds
    curve_function = CurveFitNelsonSiegelSvensson(bounds=bounds)
    fitted_curve = BondYieldCurve(from_datetime(ytm_date), bonds, ytms, curve_function)
    beta1 = curve_function._beta1
    beta2 = curve_function._beta2
    beta3 = curve_function._beta3
    beta4 = curve_function._beta4
    tau1 = curve_function._tau1
    tau2 = curve_function._tau2
    nss_stats.loc[len(nss_stats)] = [ytm_date, beta1, beta2, beta1 + beta2, beta3, beta4, tau1, tau2]
    if plot:
        today = ytm_date.strftime('%Y%m%d')
        fitted_curve.plot(f'CDB Yield Curve by Nelson Siegel Svensson @{today}')
        address = f'./results/images/nss_cdb_ytm_{today}.jpg'
        plt.savefig(address, dpi=600)
        plt.close()
    return nss_stats

def curves_stats():
    df = pd.read_csv(root_path + 'cdb_ytm.csv', parse_dates=['date'], encoding='gb2312').set_index('date')
    df['slope'] = df['10y'] - df['6m']
    df['curvature'] = df['10y'] + df['6m'] - 2 * df['2y']
    df.to_csv('./results/chinabond_ytm.csv', encoding='utf-8')
    corr_df = pd.DataFrame(columns=['maturity', 'p_1M', 'p_1Y', 'p_2.5Y'])
    for col in df.columns:
        #  autocorrelation, 1M=21 trade days, 1Y=250 trade days, 2.5Y=625 trade days
        corr_df.loc[len(corr_df)] = [col, df[col].autocorr(21), df[col].autocorr(250), df[col].autocorr(625)]
    corr_df = corr_df.set_index('maturity').round(3)
    chinabond_stats = df.describe().T.round(3)
    chinabond_stats = chinabond_stats.join(corr_df)
    chinabond_stats.to_csv('./results/chinabond_stats.csv', encoding='utf-8')
    ns_df = pd.read_csv('./results/ns_params_20100104_20220729.csv', parse_dates=['date']).set_index('date')
    b1_adfuller = adfuller(ns_df.beta1.to_numpy())
    b2_adfuller = adfuller(ns_df.beta2.to_numpy())
    b3_adfuller = adfuller(ns_df.beta3.to_numpy())
    b1_adf = b1_adfuller[0]
    b2_adf = b2_adfuller[0]
    b3_adf = b3_adfuller[0]
    b1_pval = b1_adfuller[1]
    b2_pval = b2_adfuller[1]
    b3_pval = b3_adfuller[1]
    ns_stats = ns_df[['beta1', 'beta2', 'beta3']].describe().T
    #  Beta值需要放大100倍
    for col in ns_stats.columns:
        if col is not 'count':
            ns_stats[col] = ns_stats[col] * 100
    ns_stats = ns_stats.round(3)
    corr_df = pd.DataFrame(columns=['beta', 'p_1M', 'p_1Y', 'p_2.5Y'])
    for col in ['beta1', 'beta2', 'beta3']:
        #  autocorrelation, 1M=21 trade days, 1Y=250 trade days, 2.5Y=625 trade days
        corr_df.loc[len(corr_df)] = [col, ns_df[col].autocorr(21), ns_df[col].autocorr(250), ns_df[col].autocorr(625)]
    corr_df = corr_df.set_index('beta').round(3)
    ns_stats = ns_stats.join(corr_df)
    ns_stats.at['beta1', 'ADF'] = b1_adf
    ns_stats.at['beta1', 'P_Value'] = b1_pval
    ns_stats.at['beta2', 'ADF'] = b2_adf
    ns_stats.at['beta2', 'P_Value'] = b2_pval
    ns_stats.at['beta3', 'ADF'] = b3_adf
    ns_stats.at['beta3', 'P_Value'] = b3_pval
    ns_stats.to_csv('./results/ns_params_stats.csv')

def fit_stats():
    chinabond = pd.read_csv('./results/' + 'chinabond_ytm.csv', parse_dates=['date'], encoding='gb2312').set_index(
        'date')
    ns = pd.read_csv('./results/' + 'ns_params_20100104_20220729.csv', parse_dates=['date'],
                     encoding='gb2312').set_index(
        'date')
    combined = chinabond.join(ns)
    combined = combined[chinabond.columns.values.tolist() + ['beta1', 'beta2', 'beta3', 'tau']]
    for curve in combined.itertuples():
        ns_curve = CurveFitNelsonSiegel(tau=curve.tau)
        ns_curve._beta1 = curve.beta1
        ns_curve._beta2 = curve.beta2
        ns_curve._beta3 = curve.beta3
        for chinabond_label, pred_label, ttm in zip(chinabond_ttm_labels, predicted_ttm_labels, maturities):
            combined.at[curve.Index, pred_label] = round(ns_curve._interpolated_yield(ttm) * 100, 2)
    combined.to_csv('./results/fit_stat.csv')
    real_fitted_ytm = pd.DataFrame(columns=['Indicator'] + chinabond_ttm_labels)
    real_avgs = [chinabond[t].mean() for t in chinabond_ttm_labels]
    fitted_avgs = [combined[t].mean() for t in predicted_ttm_labels]
    errors = [real - fitted for real, fitted in zip(real_avgs, fitted_avgs)]
    real_fitted_ytm.loc[0] = ['Average Chinabond YTM'] + real_avgs
    real_fitted_ytm.loc[1] = ['Average Fitted YTM'] + fitted_avgs
    real_fitted_ytm.loc[2] = ['Average YTM Errors'] + errors
    real_fitted_ytm = real_fitted_ytm.round(2).set_index('Indicator').T
    real_fitted_ytm.to_csv('./results/chinabond_vs_fitted.csv')
    corr_real_beta = combined[['beta1', 'beta2', 'beta3', '10y', 'slope', 'curvature']].corr().round(2)
    corr_real_beta.to_csv('./results/corr_real_beta.csv')
    real_fitted_ytm['maturity'] = maturities
    ax = real_fitted_ytm.plot(kind='line', x='maturity', y='Average Chinabond YTM', color='red',
                              label='Average Chinabond YTM')
    real_fitted_ytm.plot(kind='scatter', x='maturity', y='Average Fitted YTM', ax=ax, style='*',
                         label='Average Fitted YTM')
    plt.xlabel('Time to Maturity (years)')
    plt.ylabel('Yield To Maturity (%)')
    plt.grid(True)
    plt.savefig(f'./results/chinabond_vs_fitted.jpg')
    plt.close()

def plot_ytm_3d_surface():
    df = pd.read_csv(root_path + 'cdb_ytm.csv', parse_dates=['date'], encoding='gb2312')
    df = df.set_index('date').sort_index()
    x = df.columns
    y = df.index
    z = df.to_numpy()
    fig = go.Figure(data=[go.Surface(z=z, x=x, y=y)])
    fig.update_layout(
        title=dict(
            text='CDB YTM between 2010 and 2022',
            x=0.5,
            y=0.9),
        scene=dict(
            xaxis=dict(
                title='maturity',
                backgroundcolor="rgb(200, 200, 230)",
                gridcolor="white",
                showbackground=True,
                zerolinecolor="white"),
            yaxis=dict(
                title='date',
                backgroundcolor="rgb(230, 200,230)",
                gridcolor="white",
                showbackground=True,
                zerolinecolor="white"),
            zaxis=dict(
                title='ytm',
                backgroundcolor="rgb(230, 230,200)",
                gridcolor="white",
                showbackground=True,
                zerolinecolor="white"),
            aspectratio=dict(
                x=1,
                y=1,
                z=0.6
            )
        )
    )

    fig.show()

def forecast():
    factors = pd.read_csv('d:/ProjectRicequant/fxincome/eco_factors_2010_202207.csv', parse_dates=['date']).set_index(
        'date')
    factors = factors[factors.index >= datetime.datetime(2015, 12, 1)]
    ns_params = pd.read_csv('./results/fit_stat.csv', parse_dates=['date']).set_index('date')
    ns_params = ns_params[['beta1', 'beta2', 'beta3']]
    ns_params = ns_params * 100
    # pd.plotting.autocorrelation_plot(ns_params.beta3)
    X = ns_params.beta1.values
    print(X)
    train, test = X[3000: len(X) - 44], X[len(X) - 44:]
    model = AutoReg(train, lags=22)
    model_fit = model.fit()
    print(f'Coefficients: {model_fit.params}')
    predictions = model_fit.predict(start=len(train), end=len(train) + len(test) - 1, dynamic=False)
    for i in range(len(predictions)):
        print(f'predicted={predictions[i]}, expected={test[i]}')
    rmse = math.sqrt(mean_squared_error(test, predictions))
    print(f'RMSE:{rmse:.3f}')
    plt.plot(test)
    plt.plot(predictions, color='red')
    plt.show()

    ns_params = ns_params.groupby(pd.Grouper(freq='MS')).nth(10)  # 假设每个月的第11个交易日或之后才出上个月的经济数据
    ns_params = ns_params[ns_params.index >= datetime.datetime(2015, 12, 1)]
    ppi_adf = adfuller(factors.ppi.to_numpy())
    cpi_adf = adfuller(factors.cpi.to_numpy())
    m2_adf = adfuller(factors.m2.to_numpy())
    delta_adf = adfuller(factors.delta_fin_m2.to_numpy())
    # print(f'ppi adf: {ppi_adf[0]}, p: {ppi_adf[1]}')
    # print(f'cpi adf: {cpi_adf[0]}, p: {cpi_adf[1]}')
    # print(f'm2 adf: {m2_adf[0]}, p: {m2_adf[1]}')
    # print(f'delta adf: {delta_adf[0]}, p: {delta_adf[1]}')
    df = ns_params.merge(factors, on='date')
    df['beta1_lag1'] = df.beta1.shift(-1)
    df['beta2_lag1'] = df.beta2.shift(-1)
    df['beta3_lag1'] = df.beta3.shift(-1)
    df.corr().to_csv('./results/param_factors.csv', index=True)

if __name__ == '__main__':
    # preprocess()
    check_csv()
    # gen_curves('ns', start=datetime.datetime(2010, 1, 4), end=datetime.datetime(2022, 7, 29), plot=True)
    # ns_df = pd.read_csv('./results/ns_params_20100104_20220729.csv', parse_dates=['date'], encoding='gb2312')
    # plot_ns_params(ns_df)
    # curves_stats()
    # fit_stats()
    # plot_ytm_3d_surface()
    # forecast()


import logging
import socket
import sys
import zeroconf

def init_logger():
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.INFO)
    handler = logging.StreamHandler(sys.stderr)
    formatter = logging.Formatter(
        '%(asctime)s [%(levelname)s] %(message)s (at %(filename)s:%(lineno)d)')
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    return logger

class Util:

    @classmethod
    def pretty(cls, x):
        if isinstance(x, zeroconf.ServiceInfo):
            try:
                srv = x.server
                addr = socket.inet_ntoa(x.address)
                return '<ServiceInfo server=%s, address=%s>' % (srv, addr)
            except Exception:
                return str(x)
        else:
            return str(x)


import string
from enum import Enum

lower_alphabet_list = list(string.ascii_lowercase)
upper_alphabet_list = list(string.ascii_uppercase)
number_list = list(string.digits)

class EncryptionMode(Enum):
    ENC = 'ENCRYPT'
    DEC = 'DECRYPT'

def vigenere_crypt_key_range(key: str) -> int:
    key_ord = ord(key)

    if key in lower_alphabet_list:
        key_ord = key_ord - ord('a')
    elif key in upper_alphabet_list:
        key_ord = key_ord - ord('A')
    elif key in number_list:
        key_ord = key_ord - ord('0')

    return key_ord

def vigenere_crypt_pass(ch:str, key:int) -> str:
    target_list = None
    ch_ord = ord(ch)

    pos = 0
    pos_range = 26

    if ch in lower_alphabet_list:
        pos = ch_ord - ord('a')
        target_list = lower_alphabet_list
    elif ch in upper_alphabet_list:
        pos = ch_ord - ord('A')
        target_list = upper_alphabet_list
    elif ch in number_list:
        pos = ch_ord - ord('0')
        target_list = number_list
        pos_range = 10

    return target_list[(pos + key) % pos_range]

def vigenere_encrypt_decrypt(text: str, key: str, mode: EncryptionMode) -> str:
    """
    비제네르 암호를 이용하여 암호화 혹은 복호화를 수행하는 암호 알고리즘
    :param text: 암호화할 문자열
    :param key: 암호화에 사용할 key의 배열
    :param mode: 암호화할 지 복호화할 지 구분하기 위한 값
    :return: 비제네르 암호를 이용한 암호문 혹은 복호화된 문자열
    """

    key_len     = len(key)
    enc_mode    = 1 if mode == EncryptionMode.ENC else -1
    result_list = []

    for i in range(0, len(text)):
        mod_key = vigenere_crypt_key_range(key[i % key_len]) * enc_mode
        result_list.append(vigenere_crypt_pass(text[i], mod_key))

    return str(''.join(result_list))


"""Common methods we reuse."""
import collections
import os
import shutil
import stat
import tempfile as tmp
from contextlib import contextmanager
from subprocess import check_output

def sort_tag(tag):
    """Sort the list by tag."""
    list_length = len(tag)

    # Length 8 is list -l, 7 is df, 5 is list
    if list_length == 9:
        _tag = tag[4]
    elif list_length == 7:
        _tag = tag[6]
    elif list_length == 6:
        _tag = tag[3]
    else:
        _tag = tag[1]

    _sort = _tag.rsplit('_', 1)

    # We want to sort tags that have been created with count > 1. But not
    # foo_bar
    if len(_sort) > 1 and _sort[1].isdigit():
        return _sort[0], int(_sort[1])
    else:
        return _tag, 0

def sort_release(releases, split=False):
    """
    Sort the list by RELEASE, if split is true it's expecting full
    datasets.
    """
    r_dict = {}
    release_list = []

    if split:
        for rel in releases:
            rel, r_type = rel.properties["mountpoint"].value.rsplit("/")[
                -1].split("-")

            if len(rel) > 2:
                rel = float(rel)

            r_dict[rel] = r_type
    else:
        for release in releases:
            try:
                release, r_type = release.split("-")

                if len(release) > 2:
                    release = float(release)

                r_dict[release] = r_type
            except ValueError:
                pass

    ordered_r_dict = collections.OrderedDict(sorted(r_dict.items()))
    index = 0

    for r, t in ordered_r_dict.items():
        if split:
            release_list.insert(index, ["{}-{}".format(r, t)])
            index += 1
        else:
            release_list.insert(index, "{}-{}".format(r, t))
            index += 1

    return release_list

# Cyrille Pontvieux on StackOverflow
def copytree(src, dst, symlinks=False, ignore=None):
    """Copies a tree and overwrites."""
    if not os.path.exists(dst):
        os.makedirs(dst)
        shutil.copystat(src, dst)
    lst = os.listdir(src)
    if ignore:
        excl = ignore(src, lst)
        lst = [x for x in lst if x not in excl]
    for item in lst:
        s = os.path.join(src, item)
        d = os.path.join(dst, item)
        if symlinks and os.path.islink(s):
            if os.path.lexists(d):
                os.remove(d)
            os.symlink(os.readlink(s), d)
            try:
                st = os.lstat(s)
                mode = stat.S_IMODE(st.st_mode)
                os.lchmod(d, mode)
            except:
                pass  # lchmod not available
        elif os.path.isdir(s):
            copytree(s, d, symlinks, ignore)
        else:
            shutil.copy2(s, d)

# http://stackoverflow.com/questions/2333872/atomic-writing-to-file-with-python
@contextmanager
def tempfile(suffix='', dir=None):
    """
    Context for temporary file.

    Will find a free temporary filename upon entering
    and will try to delete the file on leaving, even in case of an exception.

    Parameters
    ----------
    suffix : string
        optional file suffix
    dir : string
        optional directory to save temporary file in
    """

    tf = tmp.NamedTemporaryFile(delete=False, suffix=suffix, dir=dir)
    tf.file.close()
    try:
        yield tf.name
    finally:
        try:
            os.remove(tf.name)
        except OSError as e:
            if e.errno == 2:
                pass
            else:
                raise

@contextmanager
def open_atomic(filepath, *args, **kwargs):
    """
    Open temporary file object that atomically moves to destination upon
    exiting.

    Allows reading and writing to and from the same filename.

    The file will not be moved to destination in case of an exception.

    Parameters
    ----------
    filepath : string
        the file path to be opened
    fsync : bool
        whether to force write the file to disk
    *args : mixed
        Any valid arguments for :code:`open`
    **kwargs : mixed
        Any valid keyword arguments for :code:`open`
    """
    fsync = kwargs.get('fsync', False)

    with tempfile(dir=os.path.dirname(os.path.abspath(filepath))) as tmppath:
        with open(tmppath, *args, **kwargs) as file:
            try:
                yield file
            finally:
                if fsync:
                    file.flush()
                    os.fsync(file.fileno())
        os.rename(tmppath, filepath)
        os.chmod(filepath, 0o644)

def indent_lines(message):
    """This will indent all lines except the first by 7 spaces. """
    indented = message.replace("\n", "\n       ").rstrip()

    return indented

def get_nested_key(_dict, keys=[]):
    """Gets a nested key from a dictionary."""
    key = keys.pop(0)

    if len(keys) == 0:
        return _dict[key]

    return get_nested_key(_dict[key], keys)

def checkoutput(*args, **kwargs):
    """Just a wrapper to return utf-8 from check_output"""
    out = check_output(*args, **kwargs).decode("utf-8")

    return out


import subprocess
import shlex
import os
import time
import shutil
import json

N_GROUP_CHUNKS = 4
N_EVENT_CHUNKS = 4

EXIT_FILE_MISSING = 66

# directories

def project_path():
    return os.path.join(os.path.dirname(__file__), os.path.pardir)

def to_abs(*paths):
    return os.path.join(project_path(), *paths)

def to_basename(path):
    return os.path.splitext(path)[0].rstrip("0123456789")

def to_uri(repo_name):
    return "https://github.com/christianbrugger/{}.git".format(repo_name)

# run commands

def run_secured(command, wd=project_path()):
    """ calls process and hides all outputs, usefull when operating on sensitive information """
    if isinstance(command, str):
        command = shlex.split(command)
    # do not use check=True, as it leaks secret information in tracebacks
    process = subprocess.run(command, cwd=wd,
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    if process.returncode != 0:
        raise RuntimeError("command failed")

def run(command, wd=project_path()):
    if isinstance(command, str):
        command = shlex.split(command)
    subprocess.run(command, check=True, cwd=wd)

def run_output(command, wd=project_path()):
    if isinstance(command, str):
        command = shlex.split(command)
    process = subprocess.run(command, check=True, cwd=wd,
        universal_newlines=True, stdout=subprocess.PIPE)
    return process.stdout

def run_returncode(command, wd=project_path()):
    if isinstance(command, str):
        command = shlex.split(command)
    process = subprocess.run(command, cwd=wd)
    return process.returncode

# git

def clone(repo_name):
    run("git clone {}".format(to_uri(repo_name)))

def setup_git(wd=project_path()):
    run('git config --global user.email "travis@travis-ci.org"', wd)
    run('git config --global user.name "Travis CI"', wd)

def push(repo_name, wd=project_path()):
    token = os.environ['GH_TOKEN']
    secure_uri = 'https://{}@github.com/christianbrugger/{}.git'.format(token, repo_name)

    run_secured('git remote add origin-push {}'.format(secure_uri), wd)

    while True:
        run_secured("git pull --rebase", wd)
        try:
            run_secured('git push --quiet --set-upstream origin-push master', wd)
        except RuntimeError:
            print("Error while pushing changes. Trying again in a few seconds.")
            time.sleep(3)
        else:
            break

    print("Uploaded files successfully to {}.".format(repo_name))

# parameters

def extract_parameters():
    """ extract parameters from commit log """
    message = run_output("git log -1 --pretty=%B", to_abs("..")).strip()
    data = json.loads(message)

    input_file = to_abs("..", "inputs", data["input"])
    id_ = data["id"]
    file_tag = data["tag"]
    os.environ["FILE_TAG"] = file_tag

    return input_file, id_, file_tag

# file upload

def upload_file(filenames, repo_name, id_, file_tag, do_clone=True):
    if do_clone:
        clone(repo_name)
    wd = to_abs(repo_name)
    setup_git(wd)

    if isinstance(filenames, str):
        filenames = [filenames]

    # add all files
    for filename in filenames:
        shutil.move(to_abs(filename), to_abs(repo_name, "inputs", filename))
        message = json.dumps({"input": filename, "id": id_, "tag": file_tag})
        run(['git', 'add', 'inputs/' + filename], wd)

    # check if there are changes
    files_changed = bool(run_returncode(['git', 'diff-index', '--quiet', 'HEAD', '--'], wd))

    # commit result
    if files_changed:
        run(['git', 'commit', '--message', message], wd)
        push(repo_name, wd)
    else:
        print("INFO: nothing to commit")


from typing import Any
from dataStructure.binarySearchTree import Node

class AVLNode(Node):
    def __init__(self, key: Any):
        Node.__init__(self, key)
        self.depth = 0

    def update_depth(self):
        self.depth = max(self.left_depth(), self.right_depth()) + 1

    def left_depth(self) -> int:
        return self.left.depth if self.left else 0

    def right_depth(self) -> int:
        return self.right.depth if self.right else 0

    def __str__(self) -> str:
        return str(self.key)

def ll_rotation(node: AVLNode) -> AVLNode:
    if node is None:
        return None

    left = node.left
    node.left = left.right
    left.right = node

    node.update_depth()
    left.update_depth()

    return left

def rr_rotation(node: AVLNode) -> AVLNode:
    if node is None:
        return None

    right = node.right
    node.right = right.left
    right.left = node

    node.update_depth()
    right.update_depth()

    return right

def lr_rotation(node: AVLNode) -> AVLNode:
    if node is None:
        return None

    node.left = rr_rotation(node.left)
    return ll_rotation(node)

def rl_rotation(node: AVLNode) -> AVLNode:
    if node is None:
        return None

    node.right = ll_rotation(node.right)
    return rr_rotation(node)

class AVLTree(object):
    def __init__(self):
        self.root: AVLTree = None

    def __insert_node(self, tree: AVLNode, key: Any) -> AVLNode:
        if tree is None:
            tree = AVLNode(key)
        elif key < tree.key:
            tree.left = self.__insert_node(tree.left, key)
            if tree.left_depth() - tree.right_depth() == 2:
                if key < tree.left.key:
                    tree = ll_rotation(tree)
                else:
                    tree = lr_rotation(tree)
        elif key > tree.key:
            tree.right = self.__insert_node(tree.right, key)
            if tree.right_depth() - tree.left_depth() == 2:
                if key > tree.right.key:
                    tree = rr_rotation(tree)
                else:
                    tree = rl_rotation(tree)

        tree.update_depth()
        return tree

    def insert(self, key: Any):
        if self.root is None:
            self.root = AVLNode(key)
            return
        self.root = self.__insert_node(self.root, key)

    def __delete_node(self, tree: AVLNode, z: AVLNode) -> AVLNode:
        if tree is None or z is None:
            return None

        if z.key < tree.key:
            tree.left = self.__delete_node(tree.left, z)
            if tree.right_depth() - tree.left_depth() == 2:
                r: AVLNode = tree.right
                if r.left_depth() > r.right_depth():
                    tree = rl_rotation(tree)
                else:
                    tree = rr_rotation(tree)
        elif z.key > tree.key:
            tree.right = self.__delete_node(tree.right, z)
            if tree.left_depth() - tree.right_depth() == 2:
                l: AVLNode = tree.left
                if l.right_depth() > l.left_depth():
                    tree = lr_rotation(tree)
                else:
                    tree = ll_rotation(tree)
        else:
            if tree.left and tree.right:
                if tree.left_depth() > tree.right_depth():
                    max_node = tree.left.maximum()
                    tree.key = max_node.key
                    tree.left = self.__delete_node(tree.left, max_node)
                else:
                    min_node = tree.right.minimum()
                    tree.key = min_node.key
                    tree.right = self.__delete_node(tree.right, min_node)
            else:
                tree = tree.left if tree.left else tree.right

        return tree

    def delete(self, key: Any):
        node = self.search(key)
        if node is None:
            return

        self.__delete_node(self.root, node)

    def search(self, key: Any) -> AVLNode:
        if key is None or self.root is None:
            return None

        return self.root.search(key)


# 查询集合中所有数据
# find() 方法可以查询集合中的所有数据，类似 SQL 中的 SELECT * 操作。

# 以下实例查找 sites 集合中的所有数据：

import pymongo

myclient = pymongo.MongoClient("mongodb://localhost:27017/")
mydb = myclient["runoobdb"]
mycol = mydb["sites"]

for x in mycol.find():
    print(x)


"""Tests for :mod:`siptools_research.workflow.create_logical_structmap`."""  # noqa: W505,E501

import copy
import shutil

import pytest
import lxml.etree

from metax_access import Metax

from siptools.scripts.import_description import import_description
from siptools.scripts.import_object import import_object
from siptools.scripts.compile_structmap import compile_structmap
from siptools.scripts.premis_event import premis_event
from siptools.xml.mets import NAMESPACES

import tests.utils
from siptools_research.workflow.create_logical_structmap import (
    CreateLogicalStructMap, find_dir_use_category, get_dirpath_dict
)

@pytest.mark.parametrize(
    ("events", "admids", "provenance_is_from_qvain"),
    [
        (
            [], [], False
        ),
        (
            ['creation'], ['_2005686d72f58850765d1c8147d05cb2'], False
        ),
        (
            ['creation', 'foobar'],
            ['_2005686d72f58850765d1c8147d05cb2',
             '_f8384d1f8b9cbcafcba9370d1b506a26'],
            False
        ),
        (
            ['creation'], ['_2005686d72f58850765d1c8147d05cb2'], True
        ),
    ]
)
@pytest.mark.usefixtures('testmongoclient')
def test_create_structmap_ok(
    workspace, requests_mock, events, admids, provenance_is_from_qvain
):
    """Test the workflow task CreateLogicalStructMap.

    :param workspace: Temporary workspace directory fixture
    :param requests_mock: Mocker object
    :param events: List of provenance events in dataset
    :param admids: List of identifiers of administrative metadata to
                   which the logical structuremap should refer.
    :returns: ``None``
    """
    # Create a dataset
    # Dataset contains two files
    files = [copy.deepcopy(tests.metax_data.files.BASE_FILE),
             copy.deepcopy(tests.metax_data.files.BASE_FILE)]
    files[0]['file_path'] = 'files/file1'
    files[1]['file_path'] = 'files/file2'
    dataset = copy.deepcopy(tests.metax_data.datasets.BASE_DATASET)
    # Dataset contans zero or more events
    dataset["research_dataset"]["provenance"] = []
    for event in events:
        # Provenances made in Qvain have 'lifecycle_event' instead of
        # 'provenance_event'
        if provenance_is_from_qvain:
            provenance = copy.deepcopy(
                tests.metax_data.datasets.QVAIN_PROVENANCE
            )
            provenance["lifecycle_event"]["pref_label"]["en"] = event
        else:
            provenance = copy.deepcopy(
                tests.metax_data.datasets.BASE_PROVENANCE
            )
            provenance["preservation_event"]["pref_label"]["en"] = event
        dataset["research_dataset"]["provenance"].append(provenance)
    if not dataset["research_dataset"]["provenance"]:
        del dataset["research_dataset"]["provenance"]
    tests.utils.add_metax_dataset(requests_mock, dataset=dataset, files=files)

    # Create workspace that already contains dataset files
    sip_directory = workspace / "sip-in-progress"
    sip_directory.mkdir(parents=True)
    file_directory = workspace / 'dataset_files' / 'files'
    file_directory.mkdir(parents=True)

    (file_directory / "file1").write_text("foo")
    (file_directory / "file2").write_text("bar")

    # Create required metadata to workspace:
    # * digital provenance metadata
    # * descriptive metadata
    # * technical metadata
    # * physical structure map
    for event in events:
        premis_event(
            workspace=str(sip_directory),
            event_type=event,
            event_datetime='2014-12-31T08:19:58Z',
            event_detail='foo',
            event_outcome='success',
            event_outcome_detail='bar'
        )
    if events:
        shutil.copy(
            sip_directory / 'premis-event-md-references.jsonl',
            workspace / 'create-provenance-information.jsonl'
        )
    import_description(
        dmdsec_location='tests/data/datacite_sample.xml',
        workspace=str(sip_directory)
    )
    import_object(
        workspace=str(sip_directory),
        base_path=str(workspace),
        skip_wellformed_check=True,
        filepaths=[str(file_directory)]
    )
    compile_structmap(
        workspace=str(sip_directory),
        structmap_type='Fairdata-physical'
    )

    # Init and run CreateLogicalStructMap task
    sip_prerun_files = {path.name for path in sip_directory.iterdir()}
    task = CreateLogicalStructMap(workspace=str(workspace),
                                  dataset_id='dataset_identifier',
                                  config=tests.conftest.UNIT_TEST_CONFIG_FILE)
    task.run()
    assert task.complete()

    validate_logical_structmap_file(
        str(sip_directory / 'logical_structmap.xml'), admids
    )

    sip_postrun_files = {path.name for path in sip_directory.iterdir()}

    # Nothing else should be created SIP directory
    assert sip_postrun_files \
        == sip_prerun_files | {'logical_structmap.xml'}

def test_get_dirpath_dict(requests_mock):
    """Test that get_dirpath_dict returns the correct dictionary.

    The dictionary maps dirpath to use_category.

    :param requests_mock: Mocker object
    """
    requests_mock.get(
        "https://metaksi/rest/v2/directories/1",
        json={
            "identifier": "1",
            "directory_path": "/"
        }
    )
    requests_mock.get(
        "https://metaksi/rest/v2/directories/2",
        json={
            "identifier": "2",
            "directory_path": "/test"
        }
    )

    metax_client = Metax("https://metaksi", "test", "test")
    dataset_metadata = {
        "research_dataset": {
            "directories": [
                {
                    "identifier": "1",
                    "use_category": {"pref_label": {"en": "rootdir"}}
                },
                {
                    "identifier": "2",
                    "use_category": {"pref_label": {"en": "testdir"}}
                }
            ]
        }
    }

    assert get_dirpath_dict(metax_client, dataset_metadata) == {
        "/": {"pref_label": {"en": "rootdir"}},
        "/test": {"pref_label": {"en": "testdir"}}
    }

# pylint: disable=invalid-name
def test_get_dirpath_dict_no_directories():
    """Test get_dirpath_dict function with dataset without directories.

    The function should return an empty dict when no directories are
    defined in the research_dataset.
    """
    metax_client = Metax("https://metaksi", "test", "test")
    assert not get_dirpath_dict(metax_client, {"research_dataset": {}})

def test_find_dir_use_category():
    """Test that find_dir_use_category returns the correct label."""
    dirpath_dict = {
        "/test1": {"pref_label": {"en": "testdir1"}},
        "/test2": {"pref_label": {"en": "testdir2"}}
    }
    languages = ["en"]

    # Straightforward cases
    assert find_dir_use_category("/test1", dirpath_dict, languages) \
        == "testdir1"
    assert find_dir_use_category("/test2", dirpath_dict, languages) \
        == "testdir2"

    # Closest parent that matches
    assert find_dir_use_category(
        "/test1/test", dirpath_dict, languages
    ) == "testdir1"

    # No matches
    assert not find_dir_use_category("/", dirpath_dict, languages)
    assert not find_dir_use_category("/test3", dirpath_dict, languages)

    # No directories were found in the research_dataset
    assert not find_dir_use_category("/test", {}, languages)

    # Match to root
    assert find_dir_use_category(
        "/",
        {"/": {"pref_label": {"en": "root"}}},
        languages
    ) == "root"

def validate_logical_structmap_file(logical_structmap_file, admids):
    """Validate logical structuremap XML-file.

    Checks that XML-file has the correct elements. Raises exception if
    XML is not valid.

    :param logical_structmap_file: XML file to be validated
    :param admids: List of identifiers of administrative metadata to
                   which the logical structuremap should refer.
    :returns: ``None``
    """
    tree = lxml.etree.parse(logical_structmap_file)
    assert tree.xpath('/mets:mets/mets:structMap',
                      namespaces=NAMESPACES)[0].attrib['TYPE'] \
        == "Fairdata-logical"
    assert tree.xpath('/mets:mets/mets:structMap/mets:div',
                      namespaces=NAMESPACES)[0].attrib['TYPE'] \
        == "logical"
    for admid in admids:
        assert admid in \
            tree.xpath('/mets:mets/mets:structMap/mets:div',
                       namespaces=NAMESPACES)[0].attrib['ADMID'].split()

    directories = tree.xpath(
        '/mets:mets/mets:structMap/mets:div/mets:div/@TYPE',
        namespaces=NAMESPACES
    )
    assert len(directories) == 1
    assert 'pid:urn:identifier' in directories

    assert len(tree.xpath('/mets:mets/mets:structMap/mets:div/mets:div'
                          '[@TYPE="pid:urn:identifier"]/mets:fptr/@FILEID',
                          namespaces=NAMESPACES)) == 2


from django.conf import settings
from django.forms import CheckboxSelectMultiple, RadioSelect

class HeaderSelect(RadioSelect):
    option_template_name = 'orga/widgets/header_option.html'

class MultipleLanguagesWidget(CheckboxSelectMultiple):
    option_template_name = 'orga/widgets/multi_languages_widget.html'

    def sort(self):
        self.choices = sorted(
            self.choices,
            key=lambda l: (
                (0 if l[0] in settings.LANGUAGES_OFFICIAL else 1),
                str(l[1]),
            ),
        )

    def options(self, name, value, attrs=None):
        self.sort()
        return super().options(name, value, attrs)

    def optgroups(self, name, value, attrs=None):
        self.sort()
        return super().optgroups(name, value, attrs)

    def create_option(
        self, name, value, label, selected, index, subindex=None, attrs=None
    ):
        attrs['lang'] = value
        opt = super().create_option(
            name, value, label, selected, index, subindex, attrs
        )
        opt['official'] = value in settings.LANGUAGES_OFFICIAL
        return opt


import tensorflow as tf
import numpy as np
import scipy.io as sio

from tensorflow.keras.callbacks import History
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from tensorflow.keras.utils import to_categorical

from tensorflow.keras.models import Model
from tensorflow.keras.layers import *

def MSC_1DCNN(datashape=(8192,1) , class_number = 10):
    # 一维卷积神经网络 (CNN) 模型
    input_signal = Input(shape=(8192, 1))
    # 它采用形状为(8192, 1) 的输入信号，并应用具有不同内核大小和步幅的多个并行一维卷积层
    # 多个并行一维卷积层
    x1 = Conv1D(filters=1, kernel_size=2, padding='valid', strides=2)(input_signal)
    x1 = MaxPooling1D(pool_size=256)(x1)
    x2 = Conv1D(filters=1, kernel_size=4, padding='valid', strides=4)(input_signal)
    x2 = MaxPooling1D(pool_size=128)(x2)
    x3 = Conv1D(filters=1, kernel_size=8, padding='valid', strides=8)(input_signal)
    x3 = MaxPooling1D(pool_size=64)(x3)
    x4 = Conv1D(filters=1, kernel_size=16, padding='valid', strides=16)(input_signal)
    x4 = MaxPooling1D(pool_size=32)(x4)
    x5 = Conv1D(filters=1, kernel_size=32, padding='valid', strides=32)(input_signal)
    x5 = MaxPooling1D(pool_size=16)(x5)
    x6 = Conv1D(filters=1, kernel_size=64, padding='valid', strides=64)(input_signal)
    x6 = MaxPooling1D(pool_size=8)(x6)
    x7 = Conv1D(filters=1, kernel_size=128, padding='valid', strides=128)(input_signal)
    x7 = MaxPooling1D(pool_size=4)(x7)
    x8 = Conv1D(filters=1, kernel_size=256, padding='valid', strides=256)(input_signal)
    x8 = MaxPooling1D(pool_size=2)(x8)

    xx = concatenate([x1, x2, x3, x4, x5, x6, x7, x8], axis=-2)
    # 这些层的输出沿着最后一个轴连接起来。
    xx = Flatten()(xx)
    xx = Dense(128, activation='relu')(xx)
    # 连接的输出被展平并通过具有 128个单元和ReLU激活的密集层
    output = Dense(10, activation='softmax')(xx)
    # 包含10个单元和softmax激活的密集层，它表示每个类别的输出概率
    model = Model(inputs=input_signal, outputs=output)

    return model

def normalization_processing(data):
    data_mean = data.mean()
    data_var = data.var()

    data = data - data_mean
    data = data / data_var

    return data

my_model = MSC_1DCNN(datashape=(8192, 1), class_number = 10)

datafile_path = '/content/drive/MyDrive/Colab Notebooks/all_data_DriveEnd.mat'
data = sio.loadmat(datafile_path)
category_list = ['NM','BA_007','BA_014','BA_021','OR_007','OR_014','OR_021','IR_007','IR_014','IR_021']
# category_list = ['NM','BA_014','OR_014','IR_014']

X_train = []
y_train = []
X_test1  = []
y_test1  = []
sample_length = 8192

sample_num_each_signal = 100 # in training data or testing data, each full signal actually generates 200 samples.
start_idx = np.int64(np.round(np.linspace(0,60000-sample_length,sample_num_each_signal)))
# start_idx = list(range(0,60000-sample_length+1,round(60000/sample_num_each_signal)))
for i in range(10):
    this_ctgr = category_list[i]

    for j in range(4):# dongli
        key_name = this_ctgr + '_' + str(j)
        this_ctgr_data_j = data[key_name]

        [X_train.append(this_ctgr_data_j[k:k+sample_length]) for k in start_idx]
        [y_train.append(i) for k in start_idx]
        # y_train.append(i)
        [X_test1.append(this_ctgr_data_j[k+60001:k+60001+sample_length]) for k in start_idx]
        [y_test1.append(i) for k in start_idx]
        # y_test.append(i)

print(len (X_train), len(X_test1))
X_train = np.squeeze(np.array(X_train))
X_test1 = np.squeeze(np.array(X_test1))

X_valid, X_test, y_valid, y_test = train_test_split(X_test1, y_test1, test_size=0.75, random_state=42)

X_train = np.array([normalization_processing(data) for data in X_train])
X_train = np.expand_dims(X_train,axis=2)
y_train_one_hot = to_categorical(np.array(y_train), len(category_list))

X_test = np.array([normalization_processing(data) for data in X_test])
X_test = np.expand_dims(X_test,axis=2)
y_test_one_hot = to_categorical(np.array(y_test), len(category_list))

X_valid = np.array([normalization_processing(data) for data in X_valid])
X_valid = np.expand_dims(X_valid,axis=2)
y_valid_one_hot = to_categorical(np.array(y_valid), len(category_list))

# configure training
opt = tf.keras.optimizers.Adam(lr = 0.01)
my_model.compile(optimizer=opt, loss=tf.keras.losses.CategoricalCrossentropy(), metrics=['accuracy'])

# reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, verbose=1, min_lr=0.0000001)
history_callback = History()
history = my_model.fit(
    x = X_train,
    y = y_train_one_hot ,
    batch_size = 32,
    epochs = 300,# times
    validation_data = (X_valid, y_valid_one_hot),
    shuffle=True,
    verbose = 1,
    # verbose = 0
    callbacks=[history_callback]

)

y_pred = my_model.predict(X_test) #shape: 22313*7 #y_test:22313*1
y_pred = np.argmax(y_pred, axis=1)
acc = accuracy_score(y_test, y_pred)

train_loss = history.history['loss']
train_acc = history.history['accuracy']
val_loss = history.history['val_loss']
val_acc = history.history['val_accuracy']


# Author: Bichen Wu (bichen@berkeley.edu) 08/25/2016

"""SqueezeDet Demo.

In image detection mode, for a given image, detect objects and draw bounding
boxes around them. In video detection mode, perform real-time detection on the
video stream.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cv2
import os

import numpy as np
import tensorflow as tf

from config import *
from train import _draw_box
from nets import *
import datetime

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string(
    'checkpoint', '../data_test_online/models/train/model.ckpt-49999',
    """Path to the model parameter file.""")
tf.app.flags.DEFINE_string(
    'demo_net', 'squeezeDet', """Neural net architecture.""")

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

def rotate_bound(image, angle):
    # grab the dimensions of the image and then determine the
    # center
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)

    # grab the rotation matrix (applying the negative of the
    # angle to rotate clockwise), then grab the sine and cosine
    # (i.e., the rotation components of the matrix)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))

    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY

    # perform the actual rotation and return the image
    return cv2.warpAffine(image, M, (nW, nH))

def cam_demo(outDir):
    """Detect image."""

    assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
        'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)

    with tf.Graph().as_default():
        # Load model
        if FLAGS.demo_net == 'squeezeDet':
            mc = kitti_squeezeDet_config()
            mc.BATCH_SIZE = 1
            # model parameters will be restored from checkpoint
            mc.LOAD_PRETRAINED_MODEL = False
            model = SqueezeDet(mc, FLAGS.gpu)
        elif FLAGS.demo_net == 'squeezeDet+':
            mc = kitti_squeezeDetPlus_config()
            mc.BATCH_SIZE = 1
            mc.LOAD_PRETRAINED_MODEL = False
            model = SqueezeDetPlus(mc, FLAGS.gpu)

        saver = tf.train.Saver(model.model_params)

        with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
            saver.restore(sess, FLAGS.checkpoint)
            idx_frame = 1
            # for f in sorted(glob.iglob(FLAGS.input_path)):
            cam = cv2.VideoCapture(0)
            while True:
                ret_val, img1 = cam.read()
                img1 = rotate_bound(img1, 90)
                # img1 = cv2.imread(f)
                im = img1.astype(np.float32, copy=True)

                im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
                img = cv2.resize(img1, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
                input_image = im - mc.BGR_MEANS

                # Detect
                det_boxes, det_probs, det_class = sess.run(
                    [model.det_boxes, model.det_probs, model.det_class],
                    feed_dict={model.image_input: [input_image]})

                # Filter
                final_boxes, final_probs, final_class = model.filter_prediction(
                    det_boxes[0], det_probs[0], det_class[0])

                keep_idx = [idx for idx in range(len(final_probs)) \
                            if final_probs[idx] > 0.75]  # mc.PLOT_PROB_THRESH
                final_boxes = [final_boxes[idx] for idx in keep_idx]
                final_probs = [final_probs[idx] for idx in keep_idx]
                final_class = [final_class[idx] for idx in keep_idx]

                # TODO(bichen): move this color dict to configuration file
                cls2clr = {
                    '01ball': (255, 191, 0),
                    '02basket': (0, 191, 255),
                    # 'pedestrian':(255, 0, 191)
                }

                # Draw boxes
                _draw_box(
                    img, final_boxes,
                    [mc.CLASS_NAMES[idx] + ': (%.2f)' % prob \
                     for idx, prob in zip(final_class, final_probs)],
                    cdict=cls2clr,
                )

                if not os.path.isdir(os.path.join(outDir, 'original_frames')):
                    os.makedirs(os.path.join(outDir, 'original_frames'))
                if not os.path.isdir(os.path.join(outDir, 'result_frames')):
                    os.makedirs(os.path.join(outDir, 'result_frames'))

                cv2.imwrite(os.path.join(outDir, 'original_frames', '%06d.jpg') % idx_frame, img1)
                cv2.imwrite(os.path.join(outDir, 'result_frames', '%06d.jpg') % idx_frame, img)

                idx_frame += 1
                cv2.imshow('my webcam', img)
                if cv2.waitKey(1) == 27:
                    break  # esc to quit
            cv2.destroyAllWindows()

def main(argv=None):

    now = datetime.datetime.now()
    newDirName = '../data_test_online/' + now.strftime("%Y_%m_%d-%H%M%S")

    # os.mkdir(newDirName)
    #
    # outDirs = '../data_test_online/' + today.strftime('%Y%m%d') + h
    tf.gfile.MakeDirs(newDirName)
    cam_demo(newDirName)

if __name__ == '__main__':
    tf.app.run()


import glob
import json

"""
List build files in `tools/cloud-build/daily-tests/builds/`
and creates a cron-schedule-string for each of them.
E.g.:
```
$ ./list_tests.py | jq
{
  "batch-mpi": "30 0 * * *",
  "chrome-remote-desktop": "48 0 * * *",
  "cloud-batch": "6 1 * * *",
  "hpc-high-io-v4": "24 1 * * *",
  "hpc-high-io-v5": "42 1 * * *",
...
```
"""
def list_builds(start_time=30, end_time=5*60):
    """
    start_time: int, minutes since midnight to run the first test, e.g. 30 = 00:30
    end_time: int, minutes since midnight to run the last test, e.g. 5*60 = 05:00
    """
    builds = glob.glob("*.yaml", root_dir="../daily-tests/builds/")
    assert builds, "No builds have been found"
    # Sort and strip ".yaml"
    builds = [b[:-5] for b in sorted(builds)]
    interval = (end_time - start_time) // max(1, len(builds) - 1)
    res = {}
    for b in builds:
        # OFE-deployment test is configured to only run as a PR trigger and does
        # not run on a nightly basis. Refer tools/cloud-build/provision/pr-ofe-test.tf
        # for the configuration.
        if b == "ofe-deployment":
          continue
        h, m = start_time // 60, start_time % 60
        res[b] = f"{m} {h} * * MON-FRI"
        start_time += interval
    return res

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('start_time', type=int,
                        help='minutes since midnight to run the first test, e.g. 30 = 00:30')
    parser.add_argument('end_time', type=int,
                        help='minutes since midnight to run the last test, e.g. 300 = 05:00')
    args = parser.parse_args()

    print(json.dumps(list_builds(args.start_time, args.end_time)))


from typing import List

def solve() -> None:
    with open("Input/Door_06.txt", 'r') as infile:
        answers: List[str] = infile.readlines()

    first(answers)
    second(answers)

def first(answers: List[str]) -> None:
    sum_questions: int = 0
    group_yes: List[str] = []

    for line in answers:
        if line == "\n":
            sum_questions += len(group_yes)
            group_yes = []
        else:
            for char in line:
                if char not in group_yes and char != "\n":
                    group_yes.append(char)

    sum_questions += len(group_yes)

    print("Door  6.1 | Sum of questions anyone answered yes in a group: {0}".format(sum_questions))

def second(answers: List[str]) -> None:
    sum_questions: int = 0
    group_yes: List[str] = []
    group_start: bool = True

    for line in answers:
        if line == "\n":
            sum_questions += len(group_yes)
            group_start = True
            group_yes = []
        else:
            if group_start:
                for char in line:
                    if char != "\n":
                        group_yes.append(char)

                group_start = False
            else:
                group_remove: List[str] = []
                for char in group_yes:
                    if char not in line:
                        group_remove.append(char)

                for char in group_remove:
                    group_yes.remove(char)

    sum_questions += len(group_yes)

    print("Door  6.2 | Sum of questions everyone answered yes in a group: {0}".format(sum_questions))


from django.shortcuts import render
from rest_framework.decorators import api_view
from rest_framework.response import Response
from rest_framework import status
from .models import NoteModel
from .serializers import NoteSerializer
import requests

# Create your views here.
@api_view(['GET', 'POST'])
def list(request):
    if request.method == 'GET':
        model = NoteModel.objects.all()
        serializer = NoteSerializer(model, many=True)
        return Response(data=serializer.data, status=status.HTTP_200_OK)

    if request.method == 'POST':
        serializer = NoteSerializer(data=request.data)
        if serializer.is_valid():
            serializer.save()
        return Response(serializer.data)

@api_view(['GET', 'PUT', 'DELETE'])
def detail(request, note_id):

    try:
        model = NoteModel.objects.get(id=note_id)
    except NoteModel.DoesNotExist:
        return Response(status=status.HTTP_400_BAD_REQUEST)

    if request.method == 'GET':
        serializer = NoteSerializer(model)
        return Response(data=serializer.data)

    elif request.method == 'PUT':
        serializer = NoteSerializer(model, data = request.data)
        if serializer.is_valid():
            serializer.save()
            return Response(serializer.data)
        return Response(serializer.errors, status.HTTP_400_BAD_REQUEST)

    elif request.method == 'DELETE':
        model.delete()
        return Response(status=status.HTTP_204_NO_CONTENT)

def note(request):
    notes = []

    url = "http://localhost:8000/note/api/"
    # token = "xxx"
    headers = {
        # "Authorization": f"Bearer {token}"
    }
    response = requests.get(url, headers=headers)
    notes = response.json()
    print(notes)

    context = {
        "notes": notes
    }
    return render(request, 'note/note.html', context)


import datetime
from unittest.mock import Mock
import requests

day = datetime.datetime(year=2019, month=1, day=1)
week_number = day.weekday()  # from 0 - Monday to 6 - Sunday
day_of_week = day.strftime('%A')

def is_weekday():
    today = datetime.datetime.today()
    return 0 <= today.weekday() < 5

# assert is_weekday()  # result depends on when you run the test

# create predictable resuls with mock:

# Save a couple of test days
tuesday = datetime.datetime(year=2019, month=1, day=1)
saturday = datetime.datetime(year=2019, month=1, day=5)

datetime = Mock()
# datetime.datetime.today -  is mocked method.
# use return_value  -  to mock the method.
datetime.datetime.today.return_value = tuesday
assert is_weekday()

datetime.datetime.today.return_value = saturday
assert not is_weekday()

'''
Sometimes, you’ll want to make functions return different values
when you call them more than once or even raise exceptions.
You can do this using .side_effect.

side_effect defines what happens when you call the mocked function
'''
import requests

def get_holidays():
    r = requests.get('http://localhost/api/holidays')
    if r.status_code == 200:
        return r.json()
    return None

print(get_holidays())


import flask

from flask import Flask, render_template, request, redirect, url_for
import pandas as pd
import os

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'

def validate_data(data):
    # Simplified validation: Check if "Name" column exists
    if "Name" in data.columns:
        return "Data is valid"
    else:
        return "Data is invalid"

@app.route('/', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        file = request.files['file']
        if file and file.filename.endswith('.xlsx'):
            try:
                data = pd.read_excel(file)
                validation_result = validate_data(data)
                return render_template('result.html', validation_result=validation_result)
            except Exception as e:
                return render_template('error.html', error=str(e))
        else:
            return render_template('error.html', error='Invalid file format')

    return render_template('upload.html')

if __name__ == '__main__':
    app.run(debug=True)

app.config['APPINSIGHTS_INSTRUMENTATIONKEY'] = 'a5936978-a242-45b4-ab7e-7fd4d311d24d'

appinsights = AppInsights(app)


#spark2-submit --master yarn --deploy-mode client --driver-memory 2g --num-executors 50 --executor-memory 30g --executor-cores 5 --conf spark.key1=$1 --conf spark.key2=$2 dataProfiling_spark.py
import pandas as pd
from pyspark.sql import SQLContext, SparkSession
from pyspark.sql import functions as F
from pyspark.sql.functions import isnan, when, count, col
import time

def dataprofile(data_all_df,data_cols,database_name,table_name):
    print("Starting data profiling for "+ database_name + " " + table_name + " and for "+ str(len(data_cols))+" columns")
    data_df = data_all_df.select(data_cols)

    columns2Bprofiled = data_cols
    database_name = database_name
    table_name = table_name

    dprof_df = pd.DataFrame({'database_name':[database_name] * len(data_df.columns),\
                             'table_name':[table_name] * len(data_df.columns),\
                             'column_names':data_df.columns,\
                             'data_types':[x for x in data_df.dtypes]})
    dprof_df = dprof_df[['database_name','table_name','column_names', 'data_types']]
    dprof_df.set_index('column_names', inplace=True, drop=False)

    # ======== Step 1 ==============
    num_rows = data_df.count()
    dprof_df['num_rows'] = num_rows

    print("Step 1: Num rows calculation: completed....")

    # ======== Step 2 ==============
    # number of rows with nulls and nans
    df_nacounts = data_df.select([count(when(isnan(c) | col(c).isNull(), c)).alias(c) for c in data_df.columns if data_df.select(c).dtypes[0][1]!='timestamp']).toPandas().transpose()
    df_nacounts = df_nacounts.reset_index()
    df_nacounts.columns = ['column_names','num_null']
    dprof_df = pd.merge(dprof_df, df_nacounts, on = ['column_names'], how = 'left')

    print("Step 2: Num rows with NAN/Null: completed....")

    # ============ Step 3 ============
    # number of rows with white spaces (one or more space) or blanks
    dprof_df['num_spaces'] = [data_df.where(F.col(c).rlike('^\\s+$')).count() for c in data_df.columns]
    dprof_df['num_blank'] = [data_df.where(F.col(c)=='').count() for c in data_df.columns]

    print("Step 3: White spaces/blanks: completed....")

    # ========== Step 4 ===============
    # using the in built describe() function
    desc_df = data_df.describe().toPandas().transpose()
    desc_df.columns = desc_df.iloc[0]
    desc_df=desc_df[['count', 'mean', 'stddev', 'min', 'max']]
    desc_df = desc_df.iloc[1:,:]
    desc_df = desc_df.reset_index()
    desc_df.columns.values[0] = 'column_names'
    desc_df = desc_df[['column_names','count', 'mean', 'stddev']]
    dprof_df = pd.merge(dprof_df, desc_df , on = ['column_names'], how = 'left')

    print("Step 4:Statistic metrics: completed....")

    # =============== Step 5 ============================
    # max / min values and counts
    allminvalues = [data_df.select(F.min(x)).limit(1).toPandas().iloc[0][0] for x in columns2Bprofiled]
    allmaxvalues = [data_df.select(F.max(x)).limit(1).toPandas().iloc[0][0] for x in columns2Bprofiled]

    df_counts = dprof_df[['column_names']]
    df_counts.insert(loc=0, column='min', value=allminvalues)
    df_counts.insert(loc=0, column='max', value=allmaxvalues)

    allmincounts = [data_df.where(col(x) == str(y)).count() for x, y in zip(columns2Bprofiled, allminvalues)]
    allmaxcounts = [data_df.where(col(x) == str(y)).count() for x, y in zip(columns2Bprofiled, allmaxvalues)]

    df_counts.insert(loc=0, column='counts_min', value=allmincounts)
    df_counts.insert(loc=0, column='counts_max', value=allmaxcounts)

    df_counts = df_counts[['column_names','min','counts_min','max','counts_max']]
    dprof_df = pd.merge(dprof_df, df_counts , on = ['column_names'], how = 'left')

    print("Step 5: Max/Min calculation: completed....")

    # ============= Step 6 =============================
    # number of distinct values in each column
    dprof_df['num_distinct'] = [data_df.select(x).distinct().count() for x in columns2Bprofiled]

    print("Step 6: Distinct values: completed....")

    # ================ Step 7 ============================
    # most frequently occuring value in a column and its count
    dprof_df['most_freq_valwcount'] = [data_df.groupBy(x).count().sort("count",ascending=False).limit(1).toPandas().iloc[0].values.tolist() for x in columns2Bprofiled]
    dprof_df['most_freq_value'] = [x[0] for x in dprof_df['most_freq_valwcount']]
    dprof_df['most_freq_value_count'] = [x[1] for x in dprof_df['most_freq_valwcount']]
    dprof_df = dprof_df.drop(['most_freq_valwcount'],axis=1)

    # least frequently occuring value in a column and its count
    dprof_df['least_freq_valwcount'] = [data_df.groupBy(x).count().sort("count",ascending=True).limit(1).toPandas().iloc[0].values.tolist() for x in columns2Bprofiled]
    dprof_df['least_freq_value'] = [x[0] for x in dprof_df['least_freq_valwcount']]
    dprof_df['least_freq_value_count'] = [x[1] for x in dprof_df['least_freq_valwcount']]
    dprof_df = dprof_df.drop(['least_freq_valwcount'],axis=1)

    print("Step 7: Frequency calculation: completed....")

    return dprof_df

# Set the spark session and conf parameters
spark = SparkSession.builder.appName("Data Profiling").getOrCreate()
spark.conf.set("spark.debug.maxToStringFields", 100)
spark.sparkContext.setLogLevel("ERROR")

start_time = time.time()

database_name = spark.conf.get("spark.key1")
table_name = spark.conf.get("spark.key2")
sql = 'select * from '+database_name+'.'+table_name
df = spark.sql(sql)

cols2profile = df.columns
dprofile = dataprofile(df,cols2profile,database_name,table_name)
print(dprofile)

print("\nOverall time taken..",time.time()- start_time," seconds\n")
spark.stop()


import torch
import torch.nn as nn
import numpy as np

from gensim import downloader

class IMDBNet(nn.Module):
    def __init__(self, model_args):
        super(IMDBNet, self).__init__()
        self.embedding_model = model_args.embedding_model
        self.lstm_args = model_args.lstm_args
        self.hidden_size = self.lstm_args.hidden_size if not self.lstm_args.bidirectional else self.lstm_args.hidden_size * 2
        self.output_size = model_args.output_size
        self.dropout = model_args.dropout

        self.embedding = self.load_pretrained_embedding()
        self.lstm = nn.LSTM(**self.lstm_args)
        self.classifier = nn.Sequential(
            nn.Dropout(self.dropout),
            nn.Linear(self.hidden_size, self.output_size),
            nn.Sigmoid()
        )

    def forward(self, input_ids, lengths):
        batch_size = input_ids.size(0)
        # Embed
        embeds = self.embedding(input_ids)  # (B, max_seq_length, input_size)

        # Pack and run through LSTM
        packed_embeds = torch.nn.utils.rnn.pack_padded_sequence(embeds, lengths=lengths, batch_first=True, enforce_sorted=False)
        lstm_packed_out, _ = self.lstm(packed_embeds)  # (B, max_seq_length, hidden_size)
        lstm_out, _ = torch.nn.utils.rnn.pad_packed_sequence(lstm_packed_out, batch_first=True)

        # Take the mean of the output vectors of every sequence (with consideration of their length and padding)
        seq_embeddings = (lstm_out.sum(dim=1).t() / lengths.to(lstm_out.device)).t()  # (B, hidden_size)

        # Classifier
        logits = self.classifier(seq_embeddings)  # (B, n_classes)
        logits = logits[:, 1]  # Take only the logits corresponding to 1
        logits = logits.float()
        return logits

    def load_pretrained_embedding(self):
        model = downloader.load(self.embedding_model)
        model.add_vector(len(model), np.zeros(shape=model.get_vector(0).shape))  # Unk vector
        model.add_vector(len(model) + 1, np.zeros(shape=model.get_vector(0).shape)) # Pad vector
        weights = torch.FloatTensor(model.vectors)
        embedding = nn.Embedding.from_pretrained(weights)
        return embedding

week4/unit 8/8.3 Coding Practices/modeling.py


import time

import pytest
from selenium.webdriver.common.by import By

from pages.capacitors_category_pages import film_capacitor_page
from tests import config

class TestFilmCapacitorsPage:
    @pytest.fixture()
    def film(self, driver):
        return film_capacitor_page.FilmCapacitorsPage(driver)

    # Test Film Capacitor Slick Nav
    def test_film_sub_categories(self, film):
        film.navigate_to_kemet_page()
        time.sleep(2)
        film.accept_cookies()
        film.hover_over_the_products_main_nav()
        film.hover_over_the_capacitors_sub_nav()
        film.click_on_film_category()
        time.sleep(2)
        assert ("https://www.kemet.com/en/us/capacitors/film.html" == film.get_current_url()), \
            "This is not the Film Capacitor page!"
        if film.film_category_slick_list_is_displayed:
            if film.all_film_displayed():
                print("\nFilm Capacitors Sub Categories:")
                displayed_category_tiles_elems_title = film.get_displayed_subcategory_tiles_elems_titles()
                category_tiles_elems = []
                for title in displayed_category_tiles_elems_title:
                    title_among_possible = False
                    if title.text in config.possible_film_capacitor_sub_categories:
                        title_among_possible = True
                    assert title_among_possible, "Titles don't match"
                    category_tiles_elems.append(title.text)
                    print(title.text)
            else:
                print("The All FIlm Capacitors tab is not displayed")
        else:
            print("The Slick list is not displayed")


import json

d1 = {
    'pessoa 1':{
        'nome': 'Lucas',
        'idade': 18,
    },
    'pessoa 2':{
        'nome': 'Elias',
        'idade': 22,

    }
}
d1_json = json.dumps(d1, indent=True)

with open('abc.json', 'w') as file:
    file.write(d1_json)


from __future__ import absolute_import

import json
from time import sleep

import boto

from sentry.nodestore.base import NodeStorage

def retry(attempts, func, *args, **kwargs):
    for _ in range(attempts):
        try:
            return func(*args, **kwargs)
        except Exception as err:
            sleep(0.1)
            raise
    raise

def connect_s3(bucket_name, region=None, validate=False,
               aws_access_key_id=None, aws_secret_access_key=None):
    if region is None:
        conn = boto.connect_s3(aws_access_key_id=aws_access_key_id,
                               aws_secret_access_key=aws_secret_access_key)
    else:
        conn = boto.s3.connect_to_region(region,
                                         aws_access_key_id=aws_access_key_id,
                                         aws_secret_access_key=aws_secret_access_key)
    return conn.get_bucket(bucket_name, validate=validate)

class S3NodeStorage(NodeStorage):

    def __init__(self, bucket_name=None, region=None, max_retries=3,
                 aws_access_key_id=None, aws_secret_access_key=None):
        self.max_retries = max_retries
        self.bucket = connect_s3(bucket_name, region=region,
                                 aws_access_key_id=aws_access_key_id,
                                 aws_secret_access_key=aws_secret_access_key)

    def _put(self, node_id, data):
        key = boto.s3.key.Key(self.bucket)
        key.key = node_id
        retry(self.max_retries, key.set_contents_from_string, *(data,))
        return node_id

    def delete(self, id):
        """
        >>> nodestore.delete('key1')
        """
        self.bucket.delete_key(id)

    def delete_multi(self, id_list):
        """
        Delete multiple nodes.

        Note: This is not guaranteed to be atomic and may result in a partial
        delete.

        >>> delete_multi(['key1', 'key2'])
        """
        self.bucket.delete_keys(id_list)

    def get(self, id):
        """
        >>> data = nodestore.get('key1')
        >>> print data
        """
        key = self.bucket.get_key(id)
        if key is None:
            return None
        result = retry(self.max_retries, key.get_contents_as_string)
        return json.loads(result)

    def get_multi(self, id_list):
        """
        >>> data_map = nodestore.get_multi(['key1', 'key2')
        >>> print 'key1', data_map['key1']
        >>> print 'key2', data_map['key2']
        """
        return dict(
            (id, self.get(id))
            for id in id_list
        )

    def create(self, data):
        """
        >>> nodestore.create({'foo': 'bar'})
        """
        node_id = self.generate_id()
        self._put(node_id, json.dumps(data))
        return node_id

    def set(self, id, data):
        """
        >>> nodestore.set('key1', {'foo': 'bar'})
        """
        self._put(id, json.dumps(data))

    def set_multi(self, values):
        """
        >>> nodestore.set_multi({
        >>>     'key1': {'foo': 'bar'},
        >>>     'key2': {'foo': 'baz'},
        >>> })
        """
        for id, data in values.iteritems():
            self.set(id=id, data=data)

    def cleanup(self, cutoff_timestamp):
        raise NotImplementedError


import itertools
import gzip
import csv
from collections import defaultdict

from ..common import logger
from .utils import int_or_none, named_tuple_iter, ticket_sort_key

class SenateATL:
    "parses AEC candidates Senate ATL data file (pre-2015)"

    def __init__(self, state_name, gvt_csv, firstprefs_csv):
        self.state_name = state_name
        self.gvt = defaultdict(list)
        self.ticket_votes = []
        self.btl_firstprefs = {}
        self.load_tickets(gvt_csv)
        self.load_first_preferences(state_name, firstprefs_csv)

    def load_tickets(self, gvt_csv):
        with open(gvt_csv, 'rt') as fd:
            reader = csv.reader(fd)
            # skip introduction line
            next(reader)
            header = next(reader)
            # note - this assume the GVT data is formal. FIXME: add a check for this.
            it = sorted(
                named_tuple_iter('GvtRow', reader, header, PreferenceNo=int, TicketNo=int, CandidateID=int, OwnerTicket=lambda t: t.strip()),
                key=lambda gvt: (gvt.State, ticket_sort_key(gvt.OwnerTicket), gvt.TicketNo, gvt.PreferenceNo))
            for (state_ab, ticket, ticket_no), g in itertools.groupby(
                    it, lambda gvt: (gvt.State, gvt.OwnerTicket, gvt.TicketNo)):
                if state_ab != self.state_name:
                    continue
                prefs = []
                for ticket_entry in g:
                    prefs.append(ticket_entry.CandidateID)
                self.gvt[ticket].append(tuple(prefs))

    def load_first_preferences(self, state_name, firstprefs_csv):
        with open(firstprefs_csv, 'rt') as fd:
            reader = csv.reader(fd)
            next(reader)  # skip the version
            header = next(reader)
            for idx, row in enumerate(
                named_tuple_iter(
                    'StateFirstPrefs', reader, header, TotalVotes=int, CandidateID=int)):
                if row.StateAb == state_name:
                    if row.CandidateDetails.endswith('Ticket Votes'):
                        self.ticket_votes.append((row.Ticket, row.TotalVotes))
                    elif row.CandidateDetails != 'Unapportioned':
                        assert(row.CandidateID not in self.btl_firstprefs)
                        self.btl_firstprefs[row.CandidateID] = row.TotalVotes

    def get_tickets(self):
        # GVT handling: see s272 of the electoral act (as collated in 2013)
        for group, n in self.ticket_votes:
            size = len(self.gvt[group])
            assert(size <= 3)
            remainder = n % size
            # TODO: implement some randomization here
            remainder_pattern = [1] * remainder + [0] * (size - remainder)
            remainder_pattern = [0] * (size - remainder) + [1] * (remainder)
            if remainder:
                logger.info("GVT split ticket remainder, AEO input needed: %s" % (remainder_pattern))
            # remainder_pattern = [0] * (size-remainder) + [1] * remainder
            for ticket, extra in zip(self.gvt[group], remainder_pattern):
                yield ticket, int(n / size) + extra

class SenateBTL:
    """
    build up Ticket instances for each below-the-line vote in the AEC dataset
    aggregate BTL votes with the same preference flow together
    (pre-2015)
    """

    def __init__(self, btl_csv):
        self.ticket_votes = defaultdict(int)
        self.load_btl(btl_csv)

    def load_btl(self, btl_csv):
        with gzip.open(btl_csv, 'rt') as fd:
            reader = csv.reader(fd)
            next(reader)  # skip the version
            header = next(reader)
            it = named_tuple_iter(
                'BtlRow',
                reader,
                header,
                Batch=int,
                Paper=int,
                Preference=int_or_none,
                CandidateId=int)
            self.total_ticket_data = []
            for key, g in itertools.groupby(it, lambda r: (r.Batch, r.Paper)):
                bypref = defaultdict(list)
                for row in sorted(g, key=lambda x: x.CandidateId):
                    bypref[row.Preference].append(row.CandidateId)
                flow = []
                for pref in range(1, len(bypref) + 1):
                    at_pref = bypref.get(pref)
                    if not at_pref or len(at_pref) != 1:
                        break
                    flow.append(at_pref[0])
                self.ticket_votes[tuple(flow)] += 1

    def get_tickets(self):
        for ticket in sorted(self.ticket_votes, key=lambda x: self.ticket_votes[x]):
            yield ticket, self.ticket_votes[ticket]


# !usr/bin/env python
# encoding:utf-8
import matplotlib.pyplot as plt
import torch
import cv2
import torchvision
from torchvision import transforms
import numpy as np
from mss import mss
import time

def box_iou(box1, box2):
    def box_area(box):
        return (box[2] - box[0]) * (box[3] - box[1])

    area1 = box_area(box1.T)
    area2 = box_area(box2.T)
    inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) -
             torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
    # iou = inter / (area1 + area2 - inter)
    return inter / (area1[:, None] + area2 - inter)

def xywh2xyxy(x):
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y
    return y

def xyxy2xywh(x):
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y

def plot_skeleton_kpts(im, kpts, steps, orig_shape=None):
    # 绘制coco数据集的骨架和关键点
    palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255], [153, 204, 255],
                        [255, 102, 255], [255, 51, 255], [102, 178, 255], [
                            51, 153, 255], [255, 153, 153],
                        [255, 102, 102], [255, 51, 51], [153, 255, 153], [
                            102, 255, 102], [51, 255, 51], [0, 255, 0],
                        [0, 0, 255], [255, 0, 0], [255, 255, 255]])
    skeleton = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8], [7, 9], [8, 10],
                [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]
    pose_limb_color = palette[[9, 9, 9, 9, 7, 7, 7,
                               0, 0, 0, 0, 0, 16, 16, 16, 16, 16, 16, 16]]
    pose_kpt_color = palette[[16, 16, 16, 16,
                              16, 0, 0, 0, 0, 0, 0, 9, 9, 9, 9, 9, 9]]
    radius = 9  # 节点圆点半径
    num_kpts = len(kpts) // steps  # 三个一组 51/3=17
    for kid in range(num_kpts):
        r, g, b = pose_kpt_color[kid]
        x_coord, y_coord = kpts[steps * kid], kpts[steps * kid + 1]  # 关键点坐标
        print(x_coord, y_coord)
        if not (x_coord % 640 == 0 or y_coord % 640 == 0):
            if steps == 3:
                conf = kpts[steps * kid + 2]
                if conf < 0.5:
                    continue
            cv2.circle(im, (int(x_coord), int(y_coord)),
                       radius, (int(r), int(g), int(b)), -1)
    for sk_id, sk in enumerate(skeleton):
        r, g, b = pose_limb_color[sk_id]
        pos1 = (int(kpts[(sk[0] - 1) * steps]),
                int(kpts[(sk[0] - 1) * steps + 1]))
        pos2 = (int(kpts[(sk[1] - 1) * steps]),
                int(kpts[(sk[1] - 1) * steps + 1]))
        if steps == 3:
            conf1 = kpts[(sk[0] - 1) * steps + 2]
            conf2 = kpts[(sk[1] - 1) * steps + 2]
            if conf1 < 0.5 or conf2 < 0.5:
                continue
        if pos1[0] % 640 == 0 or pos1[1] % 640 == 0 or pos1[0] < 0 or pos1[1] < 0:
            continue
        if pos2[0] % 640 == 0 or pos2[1] % 640 == 0 or pos2[0] < 0 or pos2[1] < 0:
            continue
        cv2.line(im, pos1, pos2, (int(r), int(g), int(b)), thickness=6)  # 骨架厚度

def output_to_keypoint(output):
    targets = []
    for i, o in enumerate(output):
        kpts = o[:, 6:]
        o = o[:, :6]
        for index, (*box, conf, cls) in enumerate(o.detach().cpu().numpy()):
            targets.append(
                [i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf, *list(kpts.detach().cpu().numpy()[index])])
    return np.array(targets)

def non_max_suppression_kpt(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False,
                            multi_label=False, labels=(), kpt_label=False, nc=None, nkpt=None):
    if nc is None:
        # number of classes
        nc = prediction.shape[2] - \
            5 if not kpt_label else prediction.shape[2] - 56
    xc = prediction[..., 4] > conf_thres  # candidates
    # (pixels) minimum and maximum box width and height
    min_wh, max_wh = 2, 4096
    max_det = 300  # maximum number of detections per image
    max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
    time_limit = 10.0  # seconds to quit after
    redundant = True  # require redundant detections
    multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
    merge = False  # use merge-NMS
    t = time.time()
    output = [torch.zeros((0, 6), device=prediction.device)
              ] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # image index, image inference
        x = x[xc[xi]]  # confidence
        if labels and len(labels[xi]):
            l = labels[xi]
            v = torch.zeros((len(l), nc + 5), device=x.device)
            v[:, :4] = l[:, 1:5]  # box
            v[:, 4] = 1.0  # conf
            v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
            x = torch.cat((x, v), 0)
        if not x.shape[0]:
            continue
        x[:, 5:5 + nc] *= x[:, 4:5]  # conf = obj_conf * cls_conf
        box = xywh2xyxy(x[:, :4])
        if multi_label:
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
        else:  # best class only
            if not kpt_label:
                conf, j = x[:, 5:].max(1, keepdim=True)
                x = torch.cat((box, conf, j.float()), 1)[
                    conf.view(-1) > conf_thres]
            else:
                kpts = x[:, 6:]
                conf, j = x[:, 5:6].max(1, keepdim=True)
                x = torch.cat((box, conf, j.float(), kpts), 1)[
                    conf.view(-1) > conf_thres]
        if classes is not None:
            x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
        n = x.shape[0]  # number of boxes
        if not n:  # no boxes
            continue
        elif n > max_nms:  # excess boxes
            # sort by confidence
            x = x[x[:, 4].argsort(descending=True)[:max_nms]]
        c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
        # boxes (offset by class), scores
        boxes, scores = x[:, :4] + c, x[:, 4]
        i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]
        if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
            iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
            weights = iou * scores[None]  # box weights
            x[i, :4] = torch.mm(weights, x[:, :4]).float(
            ) / weights.sum(1, keepdim=True)  # merged boxes
            if redundant:
                i = i[iou.sum(1) > 1]  # require redundancy
        output[xi] = x[i]
        if (time.time() - t) > time_limit:
            print(f'WARNING: NMS time limit {time_limit}s exceeded')
            break  # time limit exceeded
    return output

def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    shape = img.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better test mAP)
        r = min(r, 1.0)
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - \
        new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / \
            shape[0]  # width, height ratios
    dw /= 2  # divide padding into 2 sides
    dh /= 2
    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv2.copyMakeBorder(
        img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return img, ratio, (dw, dh)

# 人体姿态估计处理函数
def pose_estimation():
    screenshot = mss()  # 截图
    # img = screenshot_value.grab({"left": 960 - 200, "top": 540 - 200, "width": 400, "height": 400})
    img = screenshot.grab({"left": 1, "top": 1, "width": 800, "height": 800})
    img = cv2.cvtColor(np.array(img), cv2.COLOR_BGRA2BGR)  # 颜色空间转换(原图,格式)
    image = letterbox(img, 960, stride=64, auto=True)[0]  # 通过填充灰色 调整大小满足等比例调整
    image = transforms.ToTensor()(image)  # 转换图片到Tensor格式
    image = torch.tensor(np.array([image.numpy()]))  # 张量(数组([图片转矩阵]))
    if torch.cuda.is_available():  # 判断cuda可用
        image = image.half().to(device)  # 单精度float32转为半精度float16
    output, _ = model(image)  # 图片传入模型推演
    # 在推断结果上运行非最大抑制(NMS)(去冗余重复结果)
    output = non_max_suppression_kpt(
        output, 0.25, 0.65, nc=model.yaml['nc'], nkpt=model.yaml['nkpt'], kpt_label=True)
    with torch.no_grad():  # 不后向传播 不更新网络
        output = output_to_keypoint(output)  # 将模型输出结果转换为关键点格式
    # 图片分割 交换维度(RGB三个通道被pytorch存储tensor数据时打乱)
    nimg = image[0].permute(1, 2, 0) * 255
    nimg = nimg.cpu().numpy().astype(np.uint8)  # 转换数据类型
    nimg = cv2.cvtColor(nimg, cv2.COLOR_RGB2BGR)  # 颜色空间转换
    for idx in range(output.shape[0]):  # 遍历推断结果
        # 鼻子 - 0, 脖子 - 1，右肩 - 2，右肘 - 3，右手腕 - 4，左肩 - 5，左肘 - 6，左手腕 - 7，右臀 - 8，右膝盖 - 9，
        # 右脚踝 - 10，左臀 - 11，左膝盖 - 12，左脚踝 - 13，右眼 - 14，左眼 - 15，有耳朵 - 16，左耳朵 - 17，背景 - 18.
        # 绘制coco数据集的骨架和关键点(图片,关键点,步长)
        plot_skeleton_kpts(nimg, output[idx, 7:].T, 3)
    # pyplot进行绘图
    plt.figure(figsize=(8, 8), dpi=100, facecolor=None, edgecolor=None)  # 输出尺寸
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
    plt.margins(0, 0)
    plt.axis('off')  # 标尺线
    plt.imshow(nimg)
    plt.savefig('./output.jpg')
    # plt.show()

'''
https://github.com/WongKinYiu/yolov7
https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-w6-pose.pt
'''
# 加载已学习模型。
device = torch.device("cuda:0" if torch.cuda.is_available()
                      else "cpu")  # 识别CPU或GPU
weigths = torch.load('./yolov7-w6-pose.pt')  # 加载权重模型
model = weigths['model'].half().to(device)  # 移动所有模型参数和缓冲区到GPU
_ = model.eval()  # train()训练  eval()测试
pose_estimation()


from torch.nn.parallel import DistributedDataParallel
from torch.nn.parallel.distributed import _find_tensors
import torch.optim
import torch.utils.data
import torch
from packaging import version

class DDP(DistributedDataParallel):
    """
    Override the forward call in lightning so it goes to training and validation step respectively
    """

    def forward(self, *inputs, **kwargs):  # pragma: no cover
        if version.parse(torch.__version__[:6]) < version.parse("1.11"):
            self._sync_params()
            inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
            assert len(self.device_ids) == 1
            if self.module.training:
                output = self.module.training_step(*inputs[0], **kwargs[0])
            elif self.module.testing:
                output = self.module.test_step(*inputs[0], **kwargs[0])
            else:
                output = self.module.validation_step(*inputs[0], **kwargs[0])
            if torch.is_grad_enabled():
                # We'll return the output object verbatim since it is a freeform
                # object. We need to find any tensors in this object, though,
                # because we need to figure out which parameters were used during
                # this forward pass, to ensure we short circuit reduction for any
                # unused parameters. Only if `find_unused_parameters` is set.
                if self.find_unused_parameters:
                    self.reducer.prepare_for_backward(list(_find_tensors(output)))
                else:
                    self.reducer.prepare_for_backward([])
        else:
            from torch.nn.parallel.distributed import \
                logging, Join, _DDPSink, _tree_flatten_with_rref, _tree_unflatten_with_rref
            with torch.autograd.profiler.record_function("DistributedDataParallel.forward"):
                if torch.is_grad_enabled() and self.require_backward_grad_sync:
                    self.logger.set_runtime_stats_and_log()
                    self.num_iterations += 1
                    self.reducer.prepare_for_forward()

                # Notify the join context that this process has not joined, if
                # needed
                work = Join.notify_join_context(self)
                if work:
                    self.reducer._set_forward_pass_work_handle(
                        work, self._divide_by_initial_world_size
                    )

                # Calling _rebuild_buckets before forward compuation,
                # It may allocate new buckets before deallocating old buckets
                # inside _rebuild_buckets. To save peak memory usage,
                # call _rebuild_buckets before the peak memory usage increases
                # during forward computation.
                # This should be called only once during whole training period.
                if torch.is_grad_enabled() and self.reducer._rebuild_buckets():
                    logging.info("Reducer buckets have been rebuilt in this iteration.")
                    self._has_rebuilt_buckets = True

                # sync params according to location (before/after forward) user
                # specified as part of hook, if hook was specified.
                buffer_hook_registered = hasattr(self, 'buffer_hook')
                if self._check_sync_bufs_pre_fwd():
                    self._sync_buffers()

                if self._join_config.enable:
                    # Notify joined ranks whether they should sync in backwards pass or not.
                    self._check_global_requires_backward_grad_sync(is_joined_rank=False)

                inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
                if self.module.training:
                    output = self.module.training_step(*inputs[0], **kwargs[0])
                elif self.module.testing:
                    output = self.module.test_step(*inputs[0], **kwargs[0])
                else:
                    output = self.module.validation_step(*inputs[0], **kwargs[0])

                # sync params according to location (before/after forward) user
                # specified as part of hook, if hook was specified.
                if self._check_sync_bufs_post_fwd():
                    self._sync_buffers()

                if torch.is_grad_enabled() and self.require_backward_grad_sync:
                    self.require_forward_param_sync = True
                    # We'll return the output object verbatim since it is a freeform
                    # object. We need to find any tensors in this object, though,
                    # because we need to figure out which parameters were used during
                    # this forward pass, to ensure we short circuit reduction for any
                    # unused parameters. Only if `find_unused_parameters` is set.
                    if self.find_unused_parameters and not self.static_graph:
                        # Do not need to populate this for static graph.
                        self.reducer.prepare_for_backward(list(_find_tensors(output)))
                    else:
                        self.reducer.prepare_for_backward([])
                else:
                    self.require_forward_param_sync = False

            # TODO: DDPSink is currently enabled for unused parameter detection and
            # static graph training for first iteration.
            if (self.find_unused_parameters and not self.static_graph) or (
                    self.static_graph and self.num_iterations == 1
            ):
                state_dict = {
                    'static_graph': self.static_graph,
                    'num_iterations': self.num_iterations,
                }

                output_tensor_list, treespec, output_is_rref = _tree_flatten_with_rref(
                    output
                )
                output_placeholders = [None for _ in range(len(output_tensor_list))]
                # Do not touch tensors that have no grad_fn, which can cause issues
                # such as https://github.com/pytorch/pytorch/issues/60733
                for i, output in enumerate(output_tensor_list):
                    if torch.is_tensor(output) and output.grad_fn is None:
                        output_placeholders[i] = output

                # When find_unused_parameters=True, makes tensors which require grad
                # run through the DDPSink backward pass. When not all outputs are
                # used in loss, this makes those corresponding tensors receive
                # undefined gradient which the reducer then handles to ensure
                # param.grad field is not touched and we don't error out.
                passthrough_tensor_list = _DDPSink.apply(
                    self.reducer,
                    state_dict,
                    *output_tensor_list,
                )
                for i in range(len(output_placeholders)):
                    if output_placeholders[i] is None:
                        output_placeholders[i] = passthrough_tensor_list[i]

                # Reconstruct output data structure.
                output = _tree_unflatten_with_rref(
                    output_placeholders, treespec, output_is_rref
                )
        return output


import argparse
import json
import os
import re

from PIL import Image
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument('--input_file', required=True, type=str, help='The input SVIT JSON file')
parser.add_argument('--image_dir', required=True, type=str, help='The directory of images')
parser.add_argument('--output_file', required=True, type=str, help='The output SVIT JSON file')

def expand_to_square(box, w, h):
    if w == h:
        return box

    if w > h:
        x1, y1, x2, y2 = box
        y1 += (w - h) // 2
        y2 += (w - h) // 2
        box = x1, y1, x2, y2
        return box

    x1, y1, x2, y2 = box
    x1 += (h - w) // 2
    x2 += (h - w) // 2
    box = x1, y1, x2, y2
    return box

def normalize_bbox(box, w, h):
    if w > h:
        return [round(i / w, 3) for i in box]

    return [round(i / h, 3) for i in box]

if __name__ == "__main__":
    args = parser.parse_args()

    bbox_regex = r"\[[0-9.]+, [0-9.]+, [0-9.]+, [0-9.]+\]"
    images = dict()

    with open(args.input_file, 'r') as fin:
        data = json.load(fin)

    pbar = tqdm(data)
    for item in pbar:
        image_file_name = f"{item['image_id']}.jpg"

        if image_file_name in images:
            width, height = images[image_file_name]
        else:
            image = Image.open(os.path.join(args.image_dir, image_file_name))
            width, height = image.size
            images[image_file_name] = width, height

        for conversation in item['conversations']:
            for qa in conversation['content']:
                regex = re.compile(r'\[[0-9.]+, [0-9.]+, [0-9.]+, [0-9.]+\]')
                matches = re.findall(regex, qa['value'])
                new_value = qa['value']

                for match in matches:
                    try:
                        x1, y1, x2, y2 = [float(part.strip()) for part in match.strip('[]').split(',')]
                        original_bbox = round(x1 * width), round(y1 * height), round(x2 * width), round(y2 * height)
                        expanded_bbox = expand_to_square(original_bbox, width, height)
                        new_x1, new_y1, new_x2, new_y2 = normalize_bbox(expanded_bbox, width, height)
                        new_bbox_string = f'[{new_x1}, {new_y1}, {new_x2}, {new_y2}]'
                        new_value = new_value.replace(match, new_bbox_string)
                    except:
                        pass

                qa['value'] = new_value

    with open(args.output_file, "w+") as fout:
        fout.write(json.dumps(data, indent=4))