import os
import sys
import argparse
import time
from pathlib import Path

import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
import numpy as np
import zmq
import json
import time
import traceback

import warnings
warnings.filterwarnings('ignore')

# Constants
ROOT = Path(__file__).resolve().parents[0]
WEIGHTS = ROOT / 'weights'



from strong_sort.utils.parser import get_config
from strong_sort.strong_sort import StrongSORT
from yolov7.utils.datasets import LoadStreams, LoadImages
from yolov7.utils.general import check_imshow
from yolov7.utils.torch_utils import select_device
from yolov7.utils.plots import plot_one_box


def process_frame_and_boxes(strong, names, colors, device, frame_np, bounding_boxes):
    
    try:
        # Prepare frame
        frame = cv2.imdecode(frame_np, cv2.IMREAD_COLOR)
        im0 = frame.copy()

        # Initialize variables
        nr_sources = len(bounding_boxes)
        outputs = [None] * nr_sources
        trajectory = {}

        # Process each bounding box
        for i, bbox in enumerate(bounding_boxes):
            xmin = bbox["xmin"]
            ymin = bbox["ymin"]
            xmax = bbox["xmax"]
            ymax = bbox["ymax"]
            cls_str = bbox["class"]
            conf = bbox["probability"]
            width = xmax - xmin
            height = ymax - ymin
            #xmin, ymin, xmax, ymax, cls, conf = bbox
            
            try:
                cls = names.index(cls_str)
            except ValueError:
                print(f"Unknown class: {cls_str}")
                continue

            
            # Prepare bounding box data
            bboxes = np.array([xmin, ymin, width, height])
            #id = i  # Assuming id can be derived from the loop index
            cls = int(cls)
            conf = float(conf)
            

            # Track object using StrongSORT
            xywhs = torch.Tensor([bboxes])  # Convert to tensor
            confs = torch.Tensor([conf])
            clss = torch.Tensor([cls])
            print("BB: " ,bbox)
            
            try:
                outputs[i] = strong.update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
                print(f"Box:{xywhs} confidence:{confs} output:{outputs[i]}")
            except Exception as e:
                print(f"Unable to process properly: {e}")
            
            #outputs[i] = strongsort_list[i].update(xywhs, confs, clss, im0)
            # draw boxes for visualization
            if len(outputs[i]) > 0:
                for j, (output, conf) in enumerate(zip(outputs[i], confs)):
                    bboxes = output[0:4]
                    id = output[4]
                    cls = int(output[5])
                    # Draw trajectory if required
                    if draw:
                        center = ((int(bboxes[0]) + int(bboxes[2])) // 2, (int(bboxes[1]) + int(bboxes[3])) // 2)
                        if id not in trajectory:
                            trajectory[id] = []
                        trajectory[id].append(center)
                        for i1 in range(1, len(trajectory[id])):
                            if trajectory[id][i1 - 1] is None or trajectory[id][i1] is None:
                                continue
                            thickness = 2
                            try:
                                cv2.line(im0, trajectory[id][i1 - 1], trajectory[id][i1], (0, 0, 255), thickness)
                            except Exception as e:
                                print(f"Error drawing trajectory: {e}")
                    
                    # Draw bounding boxes on image
                    if show_vid:
                        label = f"{names[cls]} ID: {id}"
                        plot_one_box(bboxes, im0, label=label, color=colors[cls], line_thickness=line_thickness)


        # Display or save images/videos if required
            if show_vid:
                cv2.imshow("Tracking", im0)
                cv2.waitKey(1)  # Adjust as necessary for frame rate
        #if save_img:
        #   cv2.imwrite(save_path, im0)
        #if save_vid:
        #   vid_writer.write(im0)'''
    except Exception as e:
        print(f"Error: {e}")
        print(traceback.format_exc())
        
        
def set_logging():
    import logging
    logging.basicConfig(format="%(message)s", level=logging.INFO)
    
def track_bounding_boxes():
    # Process frames and boxes
    context = zmq.Context()
    socket = context.socket(zmq.SUB)
    socket.connect("tcp://127.0.0.1:5555")
    socket.setsockopt_string(zmq.SUBSCRIBE, "")
    # Initialize StrongSORT
    cfg = get_config()
    cfg.merge_from_file(opt.config_strongsort)
    
    # Placeholder for class names, should be provided
    names = ['person', 'vehicle', 'wheel', 'weapon']
    # Placeholder for colors, should be provided
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
    
    #nr_sources = 1
    # Create as many StrongSORT instances as there are sources
    #strongsort_list = []
    #for i in range(nr_sources):
     #   strongsort_list.append(
    strong = StrongSORT(
        strong_sort_weights,
        device,
        max_dist=cfg.STRONGSORT.MAX_DIST,
        max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE,
        max_age=cfg.STRONGSORT.MAX_AGE,
        n_init=cfg.STRONGSORT.N_INIT,
        nn_budget=cfg.STRONGSORT.NN_BUDGET,
        mc_lambda=cfg.STRONGSORT.MC_LAMBDA,
        ema_alpha=cfg.STRONGSORT.EMA_ALPHA,
    )
        #)

    

    try:
        while True:
            print("Listening...")
            msg = socket.recv_string()
            frame_data = socket.recv()
            frame_np = np.frombuffer(frame_data, np.uint8)
            if msg and frame_np.size > 0:
                bounding_boxes = json.loads(msg)
                print("bounding boxes: ", bounding_boxes)
                process_frame_and_boxes(strong, names, colors, device, frame_np, bounding_boxes)
            else:
                print("Error: Empty message received")
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"Error during processing: {e}")
    finally:
        socket.close()
        context.term()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--strong-sort-weights', type=str, default=WEIGHTS / 'osnet_x0_25_msmt17.pt')
    parser.add_argument('--config-strongsort', type=str, default='strong_sort/configs/strong_sort.yaml')
    parser.add_argument('--save-img', action='store_true', help='save results to *.jpg')
    parser.add_argument('--save-vid', action='store_true', help='save video tracking results')
    parser.add_argument('--show-vid', action='store_true', default=True , help='display results')
    parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
    parser.add_argument('--project', default='runs/track', help='save results to project/name')
    parser.add_argument('--exp-name', default='exp', help='save results to project/name')
    parser.add_argument('--line-thickness', default=1, type=int, help='bounding box thickness (pixels)')
    parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
    parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
    parser.add_argument('--hide-class', default=False, action='store_true', help='hide IDs')
    parser.add_argument('--draw', action='store_true', help='display object trajectory lines')
    opt = parser.parse_args()

    set_logging()


    strong_sort_weights = opt.strong_sort_weights
    config_strongsort = opt.config_strongsort
    save_img = opt.save_img
    save_vid = opt.save_vid
    show_vid = opt.show_vid
    save_txt = opt.save_txt
    line_thickness = opt.line_thickness
    draw = opt.draw
    
    save_dir = Path(opt.project) / opt.exp_name
    save_dir.mkdir(parents=True, exist_ok=True)
    
     # Initialize variables
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
      # Update as per your implementation
    
    track_bounding_boxes()

def update(self, bbox_xywh, confidences, classes, ori_img):
        try:
            print(f"Box_xywhs: {bbox_xywh} confdence:{confidences}")
            self.height, self.width = ori_img.shape[:2]
            # generate detections
            features = self._get_features(bbox_xywh, ori_img)
            
            bbox_tlwh = bbox_xywh # self._xywh_to_tlwh(bbox_xywh)
            detections = [Detection(bbox_tlwh[i], conf, features[i]) for i, conf in enumerate(
                confidences)]
            
            # run on non-maximum supression
            boxes = np.array([d.tlwh for d in detections])
            scores = np.array([d.confidence for d in detections])
            for det in detections:
                print(f"tlwh= {det.tlwh} confidence= {det.confidence}")
            print("Classes: ", classes)
            print("Confidence: ", confidences)
            
            print("Update_boxes: ", boxes)
            # update tracker
            self.tracker.predict()
            self.tracker.update(detections, classes, confidences)

            # output bbox identities
            outputs = []
            for track in self.tracker.tracks:
                if not track.is_confirmed() or track.time_since_update > 1:
                    continue

                box = track.to_tlwh()
                print("U11111111")
                x1, y1, x2, y2 = self._tlwh_to_xyxy(box)
                
                track_id = track.track_id
                class_id = track.class_id
                conf = track.conf 
                outputs.append(np.array([x1, y1, x2, y2, track_id, class_id, conf]))
                #print("Converted values:", outputs)
            if len(outputs) > 0:
                outputs = np.stack(outputs, axis=0)
            return outputs
        except Exception as e:
            print(f"Error in strongsort: {e}")
            print(traceback.format_exc())

$ python3 track_v7.py --yolo-weights yolov7.pt --strong-sort-weights osnet_x0_25_msmt17.pt --source ~/TownCenter/TownCenter.mp4 --save-vid --show-vid --device 0
Namespace(agnostic_nms=False, augment=False, classes=None, conf_thres=0.25, config_strongsort='strong_sort/configs/strong_sort.yaml', count=False, device='0', draw=False, exist_ok=False, exp_name='exp', hide_class=False, hide_conf=False, hide_labels=False, img_size=640, iou_thres=0.45, line_thickness=1, nosave=True, project='runs/track', save_conf=False, save_img=False, save_txt=False, save_vid=True, show_vid=True, source='/home/dl/TownCenter/TownCenter.mp4', strong_sort_weights='osnet_x0_25_msmt17.pt', trace=False, update=False, yolo_weights=['yolov7.pt'])
YOLOR 🚀 v0.1-115-g072f76c torch 1.13.1+cu117 CUDA:0 (NVIDIA GeForce RTX 3090, 24265.3125MB)

Traceback (most recent call last):
  File "/home/****/yolov7/StrongSORT-YOLO/yolov7/utils/google_utils.py", line 26, in attempt_download
    assets = [x['name'] for x in response['assets']]  # release assets
KeyError: 'assets'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "track_v7.py", line 387, in <module>
    detect()
  File "track_v7.py", line 90, in detect
    model = attempt_load(weights, map_location=device)  # load FP32 model
  File "/home/****/yolov7/StrongSORT-YOLO/yolov7/models/experimental.py", line 251, in attempt_load
    attempt_download(w)
  File "/home/****/yolov7/StrongSORT-YOLO/yolov7/utils/google_utils.py", line 31, in attempt_download
    tag = subprocess.check_output('git tag', shell=True).decode().split()[-1]
IndexError: list index out of range