"""
Module: Simulation
用于用户直接调用的函数，模拟某个code对应的某个噪声模型下的test circuits，并输出对应的logical error rate
"""
import multiprocessing
import os
import time
from multiprocessing import Pool
from typing import Iterable
import numpy as np
import pymatching
import stim
import stimbposd
from clifforder.Circuit.Circuit import Circuit


##  APPENDIX: 尝试导入Tesseract解码器
try:
    from tesseract_decoder import tesseract
    import tesseract_decoder
except ImportError:
    print("Tesseract decoder not found")


    class tesseract:

        def __init__(self):
            pass


        @staticmethod
        def TesseractConfig(dem, pqlimit, det_orders, det_beam=50, beam_climbing=False, no_revisit_dets=True):
            return Config(dem, det_beam)


    class Config:

        def __init__(self, dem: stim.DetectorErrorModel, det_beam):
            self.dem = dem
            self.det_beam = det_beam


        def compile_decoder(self):
            return stimbposd.bp_osd.BPOSD(model=self.dem, bp_method='min_sum', max_bp_iters=self.det_beam)

##  全局变量，用于记录当前模拟的进度
sample_number_now = multiprocessing.Value('i', 0)
error_number_now = multiprocessing.Value('i', 0)


##  CHAPTER: ====Simulation====
def Simulation(circuit: Circuit, p_noise_list: Iterable[float | int], p_measure_list: Iterable[float | int],
               allowed_variant_ratios: Iterable[float | int], decoder_name: str, decoder_params: dict | None = None, max_sample_number=100000000):
    """""        
    Tesseract decoder参数说明：
    {
    pqlimit：int对象，次数
    det_beam：int对象，分支查询数目
    beam_climbing：bool对象，是否梯度方法
    num_det_orders=num_det_orders，不知道什么阶数
    no_revisit_dets=no_revisit_dets，是否避免重复访问
    }
    ----------------------------------
    BP-OSD decoder参数说明：
    {
    bp_method：str对象，BP更新方法
    1. "product_sum"
    2. "min_sum"
    3. "min_sum_log"
    max_bp_iters：int对象，BP迭代次数，by default，50
    osd_method：str对象
        1. "osd_0"
        2. "osd_e": exhaustive OSD
        3. "osd_cs": combination-sweep OSD
        4. "osd_e"
    osd_order：int对象，OSD置换数目
    ----------------------------------
    input.circuit: 待模拟的电路
    input.p_noise_list: 噪声模型下的门噪声大小
    input.p_measure_list: 噪声模型下的测量噪声大小
    input.allowed_variant_ratios: 允许残差的相对大小
    input.decoder_name: str对象，解码器名称
    input.decoder_params: dict对象，解码器参数
    ----------------------------------
    output.p_noise_list: 噪声模型下的门噪声大小
    output.p_measure_list: 噪声模型下的测量噪声大小
    output.logical_error_rates: 逻辑错误率
    output.variants: 标准误差
    """""

    ##  PART：----输入检查----
    try:

        ##  校验噪声大小输入
        p_noise_list = [float(temp) for temp in p_noise_list]
        p_measure_list = [float(temp) for temp in p_measure_list]
        allowed_variant_ratios = [float(temp) for temp in allowed_variant_ratios]
        assert len(p_noise_list) == len(p_measure_list), "门噪声大小和测量噪声大小的数量必须相等"
        assert len(allowed_variant_ratios) >= len(p_noise_list), "允许残差的相对大小数量必须与噪声大小数量相等"

        ##  校验解码器配置
        decoder_name = str(decoder_name)
        if decoder_name in ['Tesseract', 'tesseract', 'Tesseract decoder', 'tesseract_decoder']:
            decoder_name = 'Tesseract'
        elif decoder_name in ['Matching', 'matching', 'Matching decoder', 'pymatching', 'MWPM']:
            decoder_name = 'MWPM'
        elif decoder_name in ['BP+OSD', 'bposd', 'BPOSD', 'BP+OSD decoder', 'stimbposd']:
            decoder_name = 'BP+OSD'
        else:
            raise ValueError("当前输入的解码器名称不合法")
        assert decoder_params is None or isinstance(decoder_params, dict), "解码器参数必须为字典类型"

    except Exception as e:
        raise ValueError("当前输入的模拟不合法") from e

    ##  PART：----数据预处理----
    prototype = circuit.prototype
    group_number = os.cpu_count() // 2
    logical_error_rates = []
    variants = []

    ##  PART：----遍历每一个噪声大小，模拟出logical error rate----
    for i in range(len(p_noise_list)):
        dem = prototype2dem(prototype, p_noise_list[i], p_measure_list[i])  # 获取当下噪声大小下的DEM

        ##  估算时间，选择步长
        start_time = time.time()
        task(dem, decoder_name, 20, decoder_params=decoder_params, test=True)
        end_time = time.time()
        duration = (end_time - start_time) / 100
        print('解码一个样本的时间为：', duration, '秒')
        if 50 / duration > 500000:
            sample_in_one_group_number = 500000
        elif 50 / duration < 10:
            sample_in_one_group_number = 10
        else:
            sample_in_one_group_number = int(50 / duration)

        ##  并行模拟
        lock = multiprocessing.Manager().Lock()
        global sample_number_now
        global error_number_now
        sample_number_now.value = 0
        error_number_now.value = 0
        with Pool(processes=group_number) as pool:
            results = [pool.apply_async(task, args=(dem, decoder_name, sample_in_one_group_number, lock, allowed_variant_ratios[i], decoder_params,False,
                                                    max_sample_number)) for _
                       in range(group_number)]
            _ = [result.get() for result in results]

        ##  计算逻辑错误率和标准误差
        error_rate = error_number_now.value / sample_number_now.value
        sigma = np.sqrt((error_number_now.value / sample_number_now.value - (
                error_number_now.value / sample_number_now.value) ** 2) / sample_number_now.value)
        logical_error_rates.append(error_rate)
        variants.append(sigma)

    ##  PART：----返回结果----
    return p_noise_list, p_measure_list, logical_error_rates, variants


##  CHAPTER: ====将原型转换为DEM====
def prototype2dem(prototype: list, p_noise: int | float, p_measure: int | float):
    """
    prototype: 待转换的原型
    p_noise: 门噪声大小
    p_measure: 测量噪声大小
    """

    ##  初始化
    dem_str = ''
    dem_temp = []

    ##  遍历原型中每一项，更新错误率
    for i, temp in enumerate(prototype):
        if temp[2] == 'DEP1':
            p_temp = (1 - np.sqrt(1 - 4 * p_noise / 3)) / 2
        elif temp[2] == 'DEP2':
            p_temp = (1 / 2) * (1 - (1 - 16 * p_noise / 15) ** (1 / 8))
        elif temp[2] == 'ONE_ERROR':
            p_temp = p_noise
        elif temp[2] == 'M_ERROR':
            p_temp = p_measure
        else:
            raise ValueError

        ##  合并相同的项
        flag = True
        for j, item in enumerate(dem_temp):
            if temp[1] == item[1]:
                item[0] = p_temp * (1 - item[0]) + (1 - p_temp) * item[0]
                flag = False
                break
        if flag:
            dem_temp.append([p_temp, temp[1]])

    ##  将DEM转换为文本DEM格式
    for temp in dem_temp:
        dem_str += ('\n' + f'error({temp[0]})' + temp[1])

    ##  交付stim，返回结果
    return stim.DetectorErrorModel(dem_str)


##  CHAPTER: ====用于并行的函数====
def task(dem: stim.DetectorErrorModel, decoder_name: str, number: int, lock=None, error_bar: float | int = None, decoder_params: dict = None,
         test: bool = False, max_sample_number=100000000):
    ##  生成参数下的BP+OSD decoder
    if decoder_name == 'BP+OSD':
        max_bp_iters = 50
        bp_method = 'min_sum'
        osd_method = 'osd_e'
        osd_order = 60
        if decoder_params is not None:
            if 'max_bp_iters' in decoder_params:
                max_bp_iters = decoder_params['max_bp_iters']
            if 'bp_method' in decoder_params:
                bp_method = decoder_params['bp_method']
            if 'osd_method' in decoder_params:
                osd_method = decoder_params['osd_method']
            if 'osd_order' in decoder_params:
                osd_order = decoder_params['osd_order']
        decoder = stimbposd.bp_osd.BPOSD(model=dem, bp_method=bp_method, max_bp_iters=max_bp_iters, osd_method=osd_method, osd_order=osd_order)

    ##  生成参数下的Tesseract decoder
    elif decoder_name == 'Tesseract':
        det_beam = 15
        beam_climbing = True
        no_revisit_dets = True
        num_det_orders = 16
        if decoder_params is not None:
            if 'det_beam' in decoder_params:
                det_beam = decoder_params['det_beam']
            if 'beam_climbing' in decoder_params:
                beam_climbing = decoder_params['beam_climbing']
            if 'num_det_orders' in decoder_params:
                num_det_orders = decoder_params['num_det_orders']
            if 'no_revisit_dets' in decoder_params:
                no_revisit_dets = decoder_params['no_revisit_dets']
        config = tesseract.TesseractConfig(dem=dem, pqlimit=200_000, det_beam=det_beam, beam_climbing=beam_climbing,
                                           det_orders=tesseract_decoder.utils.build_det_orders(
                                               dem=dem,
                                               num_det_orders=num_det_orders,
                                               method=tesseract_decoder.utils.DetOrder.DetIndex,
                                           ),
                                           no_revisit_dets=no_revisit_dets, )
        decoder = config.compile_decoder()

    ##  生成参数下的MWPM decoder
    elif decoder_name == 'MWPM':
        decoder = pymatching.Matching()
        decoder.from_detector_error_model(dem)
    else:
        raise NotImplementedError

    sampler = dem.compile_sampler()  # 生成采样器

    ##  如果是测试阶段，只需要执行一轮number个样本
    if test:
        detector_data, obs_data, error_data = sampler.sample(shots=number)  # 样本
        valid_detector_data = detector_data[
            np.where(np.logical_or(np.any(obs_data, axis=1), np.any(detector_data, axis=1)))[0]]  # 有效样本
        valid_obs_data = obs_data[np.where(np.logical_or(np.any(obs_data, axis=1), np.any(detector_data, axis=1)))[0]]  # 有效样本
        predictions = decoder.decode_batch(valid_detector_data)  # 解码器对每个样本的预测

        num_errors = 0
        for shot in range(len(valid_detector_data)):
            actual_for_shot = valid_obs_data[shot]
            predicted_for_shot = predictions[shot]
            if not np.array_equal(actual_for_shot, predicted_for_shot):
                num_errors += 1

    ##  正式运算需要满足误差条件才能终止
    else:
        while True:

            ##  计算样本并解码
            detector_data, obs_data, error_data = sampler.sample(shots=number)  # 生成待解码的样本
            valid_detector_data = detector_data[np.where(np.logical_or(np.any(obs_data, axis=1), np.any(detector_data, axis=1)))[0]]  # 有效样本
            valid_obs_data = obs_data[np.where(np.logical_or(np.any(obs_data, axis=1), np.any(detector_data, axis=1)))[0]]  # 有效样本
            predictions = decoder.decode_batch(valid_detector_data)  # 解码器对每个样本的预测

            ##  计算解码器错了多少次
            num_errors = 0
            for shot in range(len(valid_detector_data)):
                actual_for_shot = valid_obs_data[shot]
                predicted_for_shot = predictions[shot]
                if not np.array_equal(actual_for_shot, predicted_for_shot):
                    num_errors += 1

            ##  在进程锁下更新全局数据
            global sample_number_now
            global error_number_now
            lock.acquire()
            sample_number_now.value += number
            error_number_now.value += num_errors
            error_rate = error_number_now.value / sample_number_now.value
            sigma = np.sqrt((error_number_now.value / sample_number_now.value - (
                    error_number_now.value / sample_number_now.value) ** 2) / sample_number_now.value)
            sample_number_now_temp = sample_number_now.value
            lock.release()

            ##  设置最大的样本数
            if sigma < error_bar * error_rate or sample_number_now_temp > max_sample_number:
                break
