import os
import cv2
import math
from core_functions import read_table
from ultralytics import YOLO
import matplotlib.pyplot as plt

def analyze(model, img, name):
    # Run YOLOv8 tracking on the frame, persisting tracks between frames
    # results = model(img, conf=conf, classes=classes)
    # results = model.track(img, conf=conf, classes=classes, persist=True)
    results = model.track(img, conf=0.05, classes=[39], persist=True, max_det=1, imgsz=1600)

    # Gather the coordinates of the box
    boxes = results[0].boxes
    if len(boxes) == 0:
        print("No objects detected in " + name)
        return -1, -1

    x1, _, x2, _ = boxes.xyxyn[0].numpy()
    x = (x1 + x2) / 2
    width = x2 - x1
    
    return x, width

def interpolate(ratio):
    # N = 9
    # cal_ratio = [0.06, 0.18, 0.28, 0.39, 0.49, 0.60, 0.70, 0.82, 0.92]
    # cal_angledeg = [-25.27, -19.20, -13.46, -6.68, -0.58, 6.40, 12.37, 18.86, 24.24]
    N, cal_ratio, cal_angledeg = read_table("configs/config4.csv")
    print(cal_angledeg)
    if trace:
        print(f"Interpolate {ratio}")
    for i in range(1, N):
        ratiomin = cal_ratio[i - 1]
        ratiomax = cal_ratio[i]
        anglemin = cal_angledeg[i - 1]
        anglemax = cal_angledeg[i]
        if ratiomin <= ratio < ratiomax:
            if trace:
                print(f" i={i}, Ratio[{i - 1}]={ratiomin}, Ratio[{i}]={ratiomax}")
                print(f" i={i}, Angledeg[{i - 1}]={anglemin}, Angledeg[{i}]={anglemax}")
            deltay = anglemax - anglemin
            deltax = ratiomax - ratiomin
            slope = deltay / deltax
            if trace:
                print(f" deltax={deltax}, deltay={deltay}, slope={slope}°/m")
            delta = ratio - ratiomin
            angle = anglemin + slope * delta
            if trace:
                print(f" delta={delta}, angle={angle}°")
            return angle
    print(f"Error: unable to interpolate {ratio}")
    exit(0)


def estimate(leftd, rightd, alphar, betar, widthr, width):
    alphadeg, betadeg, thetadeg = 0.0, 0.0, 0.0
    alpharad, betarad, thetarad = 0.0, 0.0, 0.0
    alphatan, betatan, thetatan = 0.0, 0.0, 0.0
    gammadeg, gammarad, gammatan = 0.0, 0.0, 0.0
    distance = 0.0
    if trace:
        print("Estimation")
    # interpolate to get angles
    alphadeg = -interpolate(alphar)
    if trace:
        print(f" interpolate: alphar={alphar} (ratio) -> alphadeg={alphadeg}°")
    betadeg = interpolate(betar)
    if trace:
        print(f"interpolate: betar={betar} (ratio) -> betadeg={betadeg}°")
    gammadeg = interpolate(0.5 + widthr)
    if trace:
        print(f"interpolate: widthr={widthr} (ratio) -> gammadeg={gammadeg}°")
    # convert to radians
    alpharad = math.radians(alphadeg)
    if trace:
        print(f"to radians: alphadeg={alphadeg} -> alpharad={alpharad}")
    betarad = math.radians(betadeg)
    if trace:
        print(f"to radians: betadeg={betadeg} -> betarad={betarad}")
    # get tan
    alphatan = math.tan(alpharad)
    betatan = math.tan(betarad)
    if trace:
        print(f"tan(): alphatan={alphatan}, betatan={betatan}")
    # (x, y)
    y = (leftd + rightd) / (alphatan + betatan)
    x = y * alphatan - leftd
    if trace:
        print(f"position: x={x}, y={y}")
    # (distance, angle)
    distance_1 = math.sqrt(x * x + y * y)
    thetatan = x / y
    thetarad = math.atan(thetatan)
    thetadeg = math.degrees(thetarad)
    print(f"distance={distance_1}, thetatan={thetatan}, theta={thetadeg}, thetarad={thetarad}")
    # print(f"distance = {distance}")
    gammarad = math.radians(gammadeg)
    gammatan = math.tan(gammarad)
    distance_2 = width / gammatan
    print(f"distance={distance_2}, gammatan={gammatan}, gamma={gammadeg}, gammarad={gammarad}")
    return distance_1, distance_2, thetadeg


trace = False
folder_path = "imgs/"
model = YOLO('yolov8n.pt')

image_files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f)) and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif'))]

# Sort the image files to maintain a consistent order
image_files.sort()

dist = [2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6]
t_results = []

# Iterate over the images in pairs
for i in range(0, len(image_files), 2):
    # Check if there are at least two images remaining
    if i + 1 < len(image_files):
        img1_path = os.path.join(folder_path, image_files[i])
        img2_path = os.path.join(folder_path, image_files[i + 1])
        
        # Read the images
        img1 = cv2.imread(img1_path)
        img2 = cv2.imread(img2_path)
        
        # Process the images (example: print their names)
        print(f'Processing images: {image_files[i]} and {image_files[i + 1]}')

        lx, l_width = analyze(model, img1, "left")
        rx, r_width = analyze(model, img2, "right")
        width = (l_width + r_width) / 2

        results = estimate(0, 0.72, lx, rx, width, 0.08)
        print("\n\n\n")
        t_results.append(results[0]) 

        
        # (Optional) Display the images
        cv2.imshow('Image 1', img1)
        cv2.imshow('Image 2', img2)
    else:
        # If there's an odd number of images, the last one will be left without a pair
        print(f'Single image left without a pair: {image_files[i]}')

print("results")
print(t_results)

errors = [abs(dist[i]-t_results[i]) for i in range(len(dist))]
print("errors")
print(errors)

cv2.waitKey(0)  # Wait for a key press to close the images
cv2.destroyAllWindows()
# Plotting the error over distance
# plt.plot(df["reel"], df["dr1m1"], label="r1m1")
# plt.plot(df["reel"], df["dr2m1"], label="r2m1")
# plt.plot(df["reel"], df["dr3m1"], label="r3m1")
plt.plot(dist, errors)
plt.xlabel("Distance (m)")
plt.ylabel("Error (m)")
plt.title("Tracking Error over Distance")
plt.legend()
plt.show()