import os
import random
import torch
import time
import copy
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import sys
import logging
import math
import yaml
from tqdm import tqdm
from sklearn.datasets import fetch_20newsgroups
from torch.utils.data import Dataset
import torch.nn
import scipy
from scipy.optimize import minimize
from scipy.stats import wasserstein_distance
from qpsolvers import Problem, solve_problem
from torchvision.utils import save_image
from cvxopt import solvers, matrix
from torch.utils.data import Dataset
from torchvision import datasets, transforms
from argparse import Namespace, ArgumentParser
import torchvision.models
from sklearn.metrics.pairwise import rbf_kernel
from scipy.spatial.distance import pdist
from scipy.spatial import distance
from scipy.spatial.distance import cdist
from sklearn.cluster import KMeans
from collections import Counter
from sklearn.linear_model import LinearRegression
from sklearn.svm import SVC, OneClassSVM, LinearSVC, LinearSVR
from sklearn.linear_model._stochastic_gradient import SGDClassifier
from sklearn.neural_network._multilayer_perceptron import MLPClassifier
from sklearn.naive_bayes import MultinomialNB
from sklearn.tree import DecisionTreeClassifier
from sklearn.mixture import GaussianMixture
from sklearn.metrics import accuracy_score
from sklearn.calibration import CalibratedClassifierCV
# import matplotlib.pyplot as plt
from keras.preprocessing import sequence
from keras.preprocessing.text import Tokenizer
from sklearn.feature_extraction.text import TfidfVectorizer
from keras.preprocessing.sequence import pad_sequences
# from imblearn.over_sampling import SMOTE
from sentence_transformers import SentenceTransformer

class TensorDataset(Dataset):
    def __init__(self, images, labels): # images: n x c x h x w tensor
        self.images = images.detach().float()
        self.labels = labels.detach()

    def __getitem__(self, index):
        return self.images[index], self.labels[index]

    def __len__(self):
        return self.images.shape[0]

class TextDataset(Dataset):
    def __init__(self, dataset, labels):
        self.dataset = dataset.detach().int()
        self.labels = labels.detach()

    def __getitem__(self, index):
        return self.dataset[index], self.labels[index]

    def __len__(self):
        return self.dataset.shape[0]


def task2args(args, path):
    assert isinstance(path, str) and os.path.isfile(path), f"Config file {path} not found. "
    with open(path, 'r', encoding='utf-8') as file:
        raw = yaml.safe_load(file)

    def helper(raw, args):
        for key in raw:
            if isinstance(raw[key], dict):
                ret = helper(raw[key], Namespace())
                setattr(args, key, ret)
            else:
                setattr(args, key, raw[key])
        return args

    return helper(raw, args)