import numpy as np
from dataclasses import dataclass

@dataclass
class Boxes:
    limits: np.ndarray
    is_ring: bool
    index: int
    # cluster_num: List[int] = field(default_factory=list)
    is_cut_qz: bool = False
    is_cut_qxy: bool = False
    fitting_result = None


def find_box_type(is_cut_qxy, is_cut_qz, box, q, q_xy_max, q_z_max, polar_shape_0) -> bool:
    """
    Classifies a bounding box as either a 'line' or a 'peak' based on its aspect ratio.

    Args:
        box (array-like): [xmin, ymin, xmax, ymax]

    Returns:
        str: 'line' or 'peak'
    """
    result1 = True if (is_cut_qz and is_cut_qxy) else False

    xmin, ymin, xmax, ymax = box
    width = xmax - xmin
    height = ymax - ymin

    ang_hor = 0 if q < q_xy_max else np.arccos(q_xy_max/q)
    ang_vert = 0 if q < q_z_max else np.arccos(q_z_max/q)
    ratio = (np.pi/2 - ang_hor - ang_vert) / (np.pi/2)
    result3 = height > 0.6 * polar_shape_0 * ratio
    # if result3:
    #     print("found ring: ", np.rad2deg(ang_hor), np.rad2deg(ang_vert), ratio, xmin, ymin, xmax, ymax)
    return result1 or result3

    #
    # return True if shape <= height*1.5 else False




def boxes_preprocessing(detected_peaks, polar_shape, wavelength, q_abs_max,
                        q_xy_max, q_z_max):

    radius1_q = detected_peaks.radius - (detected_peaks.radius_width / 1) #2
    radius2_q = detected_peaks.radius + (detected_peaks.radius_width / 1) #2


    radius1 = np.round(radius1_q / q_abs_max * polar_shape[1]) - 1
    radius2 = np.round(radius2_q / q_abs_max * polar_shape[1]) + 1

    theta1_deg = detected_peaks.angle - (np.abs(detected_peaks.angle_width) / 2)
    theta2_deg = detected_peaks.angle + (np.abs(detected_peaks.angle_width) / 2)
    theta1 = np.round(theta1_deg / 90 * polar_shape[0]).astype(int) - 1
    theta2 = np.round(theta2_deg / 90 * polar_shape[0]).astype(int) + 1

    # to make boxes inside the img
    radius1 = np.clip(radius1, 0, polar_shape[1])
    radius2 = np.clip(radius2, 0, polar_shape[1])
    theta1 = np.clip(theta1, 0, polar_shape[0])
    theta2 = np.clip(theta2, 0, polar_shape[0])

    for i in range(len(theta1)):
        if theta1[i] == theta2[i]:
            print(theta1[i], theta2[i], detected_peaks.angle[i], detected_peaks.angle_width[i])
            raise


    boxes = np.stack([radius1,
                           theta1,
                           radius2,
                           theta2], axis=1)
    x0c = (radius1+radius2)/2
    y0c = (theta1+theta2)/2

    boxes_q_deg = np.stack([radius1_q,
                           theta1_deg,
                           radius2_q,
                           theta2_deg], axis=1)
    boxes_list = []
    for i in range(len(radius1)):
        is_cut_qxy, is_cut_qz = _get_cut_flags(radius1_q[i],theta1_deg[i],radius2_q[i],theta2_deg[i],wavelength)
        boxes_list.append(Boxes(boxes[i],
                                find_box_type(is_cut_qxy, is_cut_qz, boxes[i],
                                              detected_peaks.radius[i], q_xy_max, q_z_max, polar_shape[0]),
                                i,is_cut_qxy, is_cut_qz ))
        boxes_list[i].boxes_q_deg = boxes_q_deg[i]
        boxes_list[i].x0c = x0c[i]
        boxes_list[i].y0c = y0c[i]




    return boxes_list

def _get_cut_flags(radius1, theta1_deg, radius2, theta2_deg, wavelength):
    if theta1_deg < 2:
        is_cut_qxy = True
    else:
        is_cut_qxy = False
    is_cut_qz = _get_missing_wedge_pol(wavelength, radius2, theta2_deg)
    return is_cut_qz, is_cut_qxy

def _get_missing_wedge_pol(wavelength, q_abs, phi):
    k = 2 * np.pi / float(wavelength)
    return bool(np.abs(q_abs) > np.abs(2*k*np.cos(np.deg2rad(phi))))


def make_box_attributes(indices, boxes, fitting_result, type = None, debug = False):

    params = fitting_result['params']
    errors = fitting_result['errors']

    i = 0
    for ind in indices:
        box = boxes[ind]
        if box.is_ring and box.fitting_result is not None:
            continue
        box.fitting_result = {}
        for key in params.keys():
            if key.startswith(f'g{i}_'):
                key_changed = key.replace(f'g{i}_', '')
                box.fitting_result[key_changed] = params[key]
            elif key in ['A', 'B', 'C']:
                box.fitting_result[key] = params[key]
            elif key == 'lin_slope':
                box.fitting_result['A'] = params[key]
            elif key == 'lin_intercept':
                box.fitting_result['C'] = params[key]

        if not 'B' in box.fitting_result:
            box.fitting_result['B'] = 0
        if not 'angle' in box.fitting_result:
            box.fitting_result['angle'] = np.nan
        if not 'angle_width' in box.fitting_result:
            box.fitting_result['angle_width'] = np.inf
        if not 'theta' in box.fitting_result:
            box.fitting_result['theta'] = np.nan

        box.fitting_error = {}
        for key in errors.keys():
            if key.startswith(f'g{i}_'):
                key_changed = key.replace(f'g{i}_', '')
                box.fitting_error[key_changed] = errors[key]
            elif key in ['A', 'B', 'C']:
                box.fitting_error[key] = errors[key]
            elif key == 'lin_slope':
                box.fitting_error['A'] = errors[key]
            elif key == 'lin_intercept':
                box.fitting_error['C'] = errors[key]

        if not 'B' in box.fitting_error:
            box.fitting_error['B'] = 0
        if not 'angle' in box.fitting_error:
            box.fitting_error['angle'] = np.nan
        if not 'angle_width' in box.fitting_error:
            box.fitting_error['angle_width'] = np.inf
        if not 'theta' in box.fitting_error:
            box.fitting_error['theta'] = np.nan
        i+=1
        if debug:
            print("fitting_result",box.fitting_result)
            print("fitting_error", box.fitting_error)
