import warnings

import matplotlib

# matplotlib.use('Agg')

import networkx as nx
import matplotlib.pyplot as plt
matplotlib.use('Agg')
import numpy as np
import torch
import pandas as pd
import ruptures as rpt
from scipy.optimize import differential_evolution
from sklearn.metrics import mean_squared_error
from query_probability import query_one, load_ucr
warnings.filterwarnings('ignore')


def detect_change_points(data, model="l2", pen=1):
    algo = rpt.Pelt(model=model,min_size=1,jump=1).fit(data)
    #algo = rpt.Binseg(model=model, min_size=1, jump=1).fit(data)
    bkps = algo.predict(pen=pen)
    #print('=======',bkps)
    return bkps

def build_local_graph(data, center_index, window_size):
    """构建局部图"""
    start = max(0, center_index - window_size // 2)
    end = min(len(data), center_index + window_size // 2)
    G = nx.Graph()
    for i in range(start, end):
        for j in range(i + 1, end):
            weight = np.linalg.norm(data[i] - data[j])
            G.add_edge(i, j, weight=weight)
    return G

def find_important_nodes_via_components(G):
    """使用连通分量动态确定次要变点的数量"""
    components = list(nx.connected_components(G))
    important_nodes = []
    for component in components:
        if len(component) > 1:  # 考虑大于一个节点的组件
            important_nodes.extend(component)
    return list(important_nodes)

def process_change_points(data, change_points, window_size):
    all_important_nodes = set()
    for cp in change_points:
        G = build_local_graph(data, cp, window_size)
        important_nodes = find_important_nodes_via_components(G)
        print(cp, '------', important_nodes)
        all_important_nodes.update(important_nodes)
    all_important_nodes.update((change_points))
    return all_important_nodes

def get_magnitude(run_tag, factor, normalize,gpu):
    '''
    :param run_tag:
    :param factor:
    :return: Perturbed Magnitude
    '''
    data = load_ucr('data/' + run_tag + '/' + run_tag + '_attack'+gpu+'.txt', normalize=normalize)
    X = data[:, 1:]

    max_magnitude = X.max(1)
    min_magnitude = X.min(1)
    mean_magnitude = np.mean(max_magnitude - min_magnitude)

    perturbed_mag = mean_magnitude * factor
    print('Perturbed Magnitude:', perturbed_mag)

    return perturbed_mag


class Attacker:

    def __init__(self, run_tag, model_type, cuda, normalize, e,device,gpu):
        self.run_tag = run_tag
        self.model_type = model_type
        self.cuda = cuda
        #self.intervals = get_attack_position(self.run_tag, self.top_k)
        self.normalize = normalize
        self.e = e
        self.device = device
        self.gpu = gpu

    def perturb_ts(self, perturbations, ts, attack_pos):
        '''
        :param perturbations:formalized as a tuple（x,e),x(int) is the x-coordinate，e(float) is the epsilon,e.g.,(2,0.01)
        :param ts: time series
        :return: perturbed ts
        '''
        # first we copy a ts

        ts_tmp = np.copy(ts)

        coordinate = 0  # 初始化perturbations数组的索引
        for i in range(len(attack_pos)):
            if attack_pos[i] == 1:
                ts_tmp[i] += perturbations[coordinate]
                coordinate += 1
        # for interval in self.intervals:
        #     for i in range(int(interval[0]), int(interval[1])):
        #         ts_tmp[i] += perturbations[coordinate]
        #         coordinate += 1
        return ts_tmp

    def plot_per(self, perturbations, ts, target_class, sample_idx,attack_pos, prior_probs, attack_probs, factor):

        # Obtain the perturbed ts
        ts_tmp = np.copy(ts)
        ts_perturbed = self.perturb_ts(perturbations=perturbations, ts=ts, attack_pos=attack_pos)
        # Start to plot
        plt.figure(figsize=(6, 4))
        plt.plot(ts_tmp, color='b', label='Original %.2f' % prior_probs)
        plt.plot(ts_perturbed, color='r', label='Perturbed %.2f' % attack_probs)
        plt.xlabel('Time', fontsize=12)

        if target_class == -1:
            plt.title('Untargeted: Sample %d, eps_factor=%.3f' %
                      (sample_idx, factor), fontsize=14)
        else:
            plt.title('Targeted(%d): Sample %d, eps_factor=%.3f' %
                      (target_class, sample_idx, factor), fontsize=14)

        plt.legend(loc='upper right', fontsize=8)
        plt.savefig('result_' + str(factor) + '_' + str(self.model_type) + '/'
                    + self.run_tag + '/figures'+self.gpu+'/' + self.run_tag +'_' + str(sample_idx) + '.png')
        # plt.show()

    def fitness(self, device,perturbations, ts, sample_idx, queries,attack_pos, target_class=-1):
        #device = torch.device("cuda:0" if self.cuda else "cpu")
        perturbations = torch.tensor(perturbations).to(device)
        #ts = torch.tensor(ts, device='cuda')

        queries[0] += 1

        ts_perturbed = self.perturb_ts(perturbations, ts,attack_pos = attack_pos)
        prob, _, _, _, _ = query_one(run_tag=self.run_tag,device=device, idx=sample_idx, attack_ts=ts_perturbed,
                                     target_class=target_class, normalize=self.normalize,
                                     cuda=self.cuda, model_type=self.model_type, e=self.e,gpu=self.gpu)
        prob = torch.tensor(prob)

        if target_class != -1:
            prob = 1 - prob

        return prob  # The fitness function is to minimize the fitness value

    def attack_success(self,device,perturbations, ts, sample_idx, attack_pos, iterations, target_class=-1, verbose=True):
        iterations[0] += 1
        print('The %d iteration' % iterations[0])
        ts_perturbed = self.perturb_ts(perturbations, ts, attack_pos)
        #ts_perturbed = torch.tensor(ts_perturbed, device='cuda')
        # Obtain the perturbed probability vector and the prior probability vector
        prob, prob_vector, prior_prob, prior_prob_vec, real_label = query_one(self.run_tag, device,idx=sample_idx,
                                                                              attack_ts=ts_perturbed,
                                                                              target_class=target_class,
                                                                              normalize=self.normalize,
                                                                              verbose=verbose, cuda=self.cuda,
                                                                              model_type=self.model_type,
                                                                              e=self.e,gpu=self.gpu)

        predict_class = torch.argmax(prob_vector).to(device)
        prior_class = torch.argmax(prior_prob_vec).to(device)
        real_label = real_label.to(device)

        # Conditions for early termination(empirical-based estimation), leading to save the attacking time
        # But it may judge incorrectly that this may decrease the success rate of the attack.
        if (iterations[0] > 5 and prob > 0.99) or \
                (iterations[0] > 20 and prob > 0.9):
            print('The %d sample is not expected to successfully attack.' % sample_idx)
            print('prob: ',prob)
            return True

        if prior_class != real_label:
            print('The %d sample cannot be classified correctly, no need to attack' % sample_idx)
            return True

        if prior_class == target_class:
            print(
                'The true label of %d sample equals to target label, no need to attack' % sample_idx)
            return True

        if verbose:
            print('The Confidence of current iteration: %.4f' % prob)
            print('########################################################')

        # The criterion of attacking successfully:
        # Untargeted attack: predicted label is not equal to the original label.
        # Targeted attack: predicted label is equal to the target label.
        if ((target_class == -1 and predict_class != prior_class) or
                (target_class != -1 and predict_class == target_class)):
            print('##################### Attack Successfully! ##########################')

            return True

    def attack(self, sample_idx,device, target_class=-1, factor=0.04,
               max_iteration=50, popsize=200, verbose=True):

        test = load_ucr('data/' + self.run_tag + '/' + self.run_tag + '_attack'+self.gpu+'.txt'
                        , normalize=self.normalize)
        #attack_poses = np.loadtxt('data/' + self.run_tag + '/' + self.run_tag + '_attackPos'+self.gpu+'.txt')

        ori_ts = test[sample_idx][1:]
        #变点检测找区间
        length = ori_ts.shape

        result = detect_change_points(ori_ts)
        sampled_indices = np.zeros(len(ori_ts), dtype=int)
        if length[0] > 500 or len(result) > 0.4 * length[0]:

            # 设置基本窗口大小和基本采样密度
            base_window_size = 0
            base_density_factor = 0.01
        else:
            # 设置基本窗口大小和基本采样密度
            base_window_size = int(0.1 * length[0])
            base_density_factor = 0.5

        # 计算变点的重要性并调整窗口大小和采样密度
        for i, cp in enumerate(result[:-1]):
            # 计算变点前后的统计差异
            if i == 0:
                pre_mean = np.mean(ori_ts[:cp])
            else:
                pre_mean = np.mean(ori_ts[result[i - 1]:cp])
            post_mean = np.mean(ori_ts[cp:result[i + 1]])

            # 重要性评估：差异的绝对值
            importance = abs(post_mean - pre_mean)

            # 动态调整窗口大小和采样密度
            if length[0] > 500 or len(result) > 0.1 * length[0]:
                window_size = int(base_window_size + int(importance * 5)) # 重要性越大，窗口越大
                density_factor = base_density_factor + importance * 0.1  # 重要性越大，采样密度越高
            else:
                window_size = int(base_window_size + int(importance * 10))  # 重要性越大，窗口越大
                density_factor = base_density_factor + importance * 0.3  # 重要性越大，采样密度越高


            # 设置窗口的开始和结束位置
            start = max(cp - window_size // 2, 0)
            end = min(cp + window_size // 2, len(ori_ts))

            # 在窗口内随机采样
            sample_count = int((end - start) * density_factor)
            if sample_count > 0:
                samples = np.random.choice(range(start, end), sample_count, replace=True)
                sampled_indices[samples] = 1
        attack_pos = sampled_indices

        #attack_pos = attack_poses[sample_idx,:]
        # attack_pos = attack_poses.head(sample_idx)
        # attack_pos = attack_pos.apply(pd.to_numeric, errors='coerce').fillna(0)

        #ori_ts = torch.tensor(ori_ts).to(device)
        attacked_probs, attacked_vec, prior_probs, prior_vec, real_label = query_one(self.run_tag,device,idx=sample_idx,
                                                                                     attack_ts=ori_ts,
                                                                                     target_class=target_class,
                                                                                     normalize=self.normalize,
                                                                                     verbose=False,
                                                                                     cuda=self.cuda, e=self.e,
                                                                                     model_type=self.model_type,gpu=self.gpu)
        prior_class = torch.argmax(prior_vec).to(device)
        if prior_class != real_label:
            print('The %d sample cannot be classified correctly, no need to attack' % sample_idx)
            return ori_ts,ori_ts, [prior_probs, attacked_probs, 0, 0, 0, 0, 0, 'WrongSample']


        # Get the maximum perturbed magnitude
        perturbed_magnitude = get_magnitude(self.run_tag, factor, normalize=self.normalize,gpu=self.gpu)

        bounds = []

        # for interval in self.intervals:
        #     steps_count += int(interval[1]) - int(interval[0])
        #     for i in range(int(interval[0]), int(interval[1])):
        #         bounds.append((-1 * perturbed_magnitude, perturbed_magnitude))
        steps_count = attack_pos.sum()
        for i in range(len(attack_pos)):
            if attack_pos[i] == 1:
                bounds.append((-1 * perturbed_magnitude, perturbed_magnitude))

        print('The length of shapelet interval', steps_count)
        #print('The length of bounds', len(bounds))
        popmul = max(1, popsize // len(bounds))
        # Record of the number of iterations
        iterations = [0]
        queries = [0]

        def fitness_fn(perturbations):

            return self.fitness(perturbations=perturbations, ts=ori_ts, queries=queries,
                                sample_idx=sample_idx, attack_pos=attack_pos,target_class=target_class,device=device)

        def callback_fn(x, convergence):

            return self.attack_success(perturbations=x, ts=ori_ts,
                                       sample_idx=sample_idx,
                                       attack_pos = attack_pos,
                                       iterations=iterations,
                                       target_class=target_class,
                                       verbose=verbose,device=device)

        attack_result = differential_evolution(func=fitness_fn, bounds=bounds
                                               , maxiter=max_iteration, popsize=popmul
                                               , recombination=0.7, callback=callback_fn,
                                               atol=-1, polish=False)

        attack_ts = self.perturb_ts(attack_result.x, ori_ts, attack_pos)

        mse = mean_squared_error(ori_ts, attack_ts)

        attacked_probs, attacked_vec, prior_probs, prior_vec, real_label = query_one(self.run_tag,device, idx=sample_idx,
                                                                                     attack_ts=attack_ts,
                                                                                     target_class=target_class,
                                                                                     normalize=self.normalize,
                                                                                     verbose=False,
                                                                                     cuda=self.cuda, e=self.e,
                                                                                     model_type=self.model_type,gpu=self.gpu)

        predicted_class = torch.argmax(attacked_vec).to(device)
        prior_class = torch.argmax(prior_vec).to(device)

        if prior_class != real_label:
            success = 'WrongSample'

        elif prior_class == target_class:
            success = 'NoNeedAttack'

        else:
            if (predicted_class.item() != prior_class.item() and target_class == -1) \
                    or (predicted_class.item() == target_class and target_class != -1):
                success = 'Success'
            else:
                success = 'Fail'

        if success == 'Success':
            self.plot_per(perturbations=attack_result.x, ts=ori_ts, target_class=target_class,
                          sample_idx=sample_idx, attack_pos=attack_pos, prior_probs=prior_probs, attack_probs=attacked_probs, factor=factor)

        return ori_ts, attack_ts, [prior_probs, attacked_probs, prior_class.item(),
                           predicted_class.item(), queries[0], mse, iterations[0], success]


if __name__ == '__main__':
    idx = [1]
    attacker = Attacker(run_tag='ECG200', top_k=3, model_type='f', e=1499,
                        cuda=True, normalize=False)
    for idx in idx:
        attack_ts, info = attacker.attack(sample_idx=idx, target_class=-1, factor=0.01,
                                          max_iteration=200, popsize=1, verbose=True)
        #这里是生成对抗样本
        if info[-1] == 'Success':
            file = open('attack_time_series.txt', 'w+')
            file.write('%d %d ' % (idx, info[3]))  # save the sample index and the perturbed label
            for i in attack_ts:
                file.write('%.4f ' % i)
            file.write('\n')
            file.close()

        print(info)
        file = open('info.txt', 'w+')
        file.write('%d ' % idx)
        for i in info:
            if isinstance(i, int):
                file.write('%d ' % i)
            elif isinstance(i, float):
                file.write('%.4f ' % i)
            else:
                file.write(i + ' ')

        file.write('\n')
        file.close()