
import json
import os
import urllib

import cv2
import numpy as np
# import altair as alt
import pandas as pd
import streamlit as st
from jaitool.inference import D2Inferer
from jaitool.draw import draw_bbox
from seaborn import color_palette

# import streamlit_theme as stt

# Streamlit encourages well-structured code, like starting execution in a main() function.
def main():
    st.title('Electrical Drawing Analyser')
    # Render the readme as markdown using st.markdown.
    readme_text = st.markdown(get_file_content_as_string("instructions.md"))

    # Download external dependencies.
    for filename in EXTERNAL_DEPENDENCIES.keys():
        download_file(filename)

    # Once we have the dependencies, add a selector for the app mode on the sidebar.
    # st.sidebar.title("What to do")
    # app_mode = st.sidebar.selectbox("Choose the app mode",
    #     ["Show instructions", "Run the app", "Show the source code"])
    # if app_mode == "Show instructions":
    #     st.sidebar.success('To continue select "Run the app".')
    # elif app_mode == "Show the source code":
    #     readme_text.empty()
    #     st.code(get_file_content_as_string("streamlit_app.py"))
    # elif app_mode == "Run the app":
    #     readme_text.empty()
    #     run_the_app()
    readme_text.empty()
    run_the_app()

# This file downloader demonstrates Streamlit animation.
def download_file(file_path):
    # Don't download the file twice. (If possible, verify the download using the file length.)
    if os.path.exists(file_path):
        if "size" not in EXTERNAL_DEPENDENCIES[file_path]:
            return
        elif os.path.getsize(file_path) == EXTERNAL_DEPENDENCIES[file_path]["size"]:
            return

    # These are handles to two visual elements to animate.
    weights_warning, progress_bar = None, None
    try:
        weights_warning = st.warning("Downloading %s..." % file_path)
        progress_bar = st.progress(0)
        with open(file_path, "wb") as output_file:
            with urllib.request.urlopen(EXTERNAL_DEPENDENCIES[file_path]["url"]) as response:
                length = int(response.info()["Content-Length"])
                counter = 0.0
                MEGABYTES = 2.0 ** 20.0
                while True:
                    data = response.read(8192)
                    if not data:
                        break
                    counter += len(data)
                    output_file.write(data)

                    # We perform animation by overwriting the elements.
                    weights_warning.warning("Downloading %s... (%6.2f/%6.2f MB)" %
                        (file_path, counter / MEGABYTES, length / MEGABYTES))
                    progress_bar.progress(min(counter / length, 1.0))

    # Finally, we remove these visual elements by calling .empty().
    finally:
        if weights_warning is not None:
            weights_warning.empty()
        if progress_bar is not None:
            progress_bar.empty()

# This is the main app app itself, which appears when the user selects "Run the app".
def run_the_app():
    # To make Streamlit fast, st.cache allows us to reuse computation across runs.
    # In this common pattern, we download data from an endpoint only once.
    @st.cache
    def load_metadata(url):
        return pd.read_csv(url)

    # This function uses some Pandas magic to summarize the metadata Dataframe.
    @st.cache
    def create_summary(metadata):
        one_hot_encoded = pd.get_dummies(metadata[["frame", "label"]], columns=["label"])
        summary = one_hot_encoded.groupby(["frame"]).sum().rename(columns={
            "label_biker": "biker",
            "label_car": "car",
            "label_pedestrian": "pedestrian",
            "label_trafficLight": "traffic light",
            "label_truck": "truck"
        })
        return summary

    # An amazing property of st.cached functions is that you can pipe them into
    # one another to form a computation DAG (directed acyclic graph). Streamlit
    # recomputes only whatever subset is required to get the right answer!
    metadata = load_metadata(os.path.join(DATA_URL_ROOT, "labels.csv.gz"))
    summary = create_summary(metadata)

    # Uncomment these lines to peek at these DataFrames.
    # st.write('## Metadata', metadata[:1000], '## Summary', summary[:1000])

    # Draw the UI elements to search for objects (pedestrians, cars, etc.)
    selected_frame_index, selected_frame  = frame_selector_ui(summary)
    if selected_frame_index == None:
        st.error("No frames fit the criteria. Please select different label or number.")
        return

    # Draw the UI element to select parameters for the YOLO object detector.
    confidence_threshold, overlap_threshold = object_detector_ui()
    selected_classes = class_selector_ui()
    # Load the image from S3.
    image_url = os.path.join(IMG_DATA_PATH, selected_frame)
    image = load_image(image_url)

    # Add boxes for objects on the image. These are the boxes for the ground image.
    boxes = metadata[metadata.frame == selected_frame].drop(columns=["frame"])
    # draw_image_with_boxes2(image, boxes, "Ground Truth",
    #     f"**Human-annotated data** (frame `{selected_frame_index}`, `{IMG_NAMES[selected_frame_index]}`)", confidence_threshold)

    # Get the boxes for the objects detected by YOLO by running the YOLO model.
    # yolo_boxes = yolo_v3(image, confidence_threshold, overlap_threshold)
    predict_dict = rcnn(image, confidence_threshold, overlap_threshold)
    draw_image_with_boxes2(image, predict_dict, "Real-time Computer Vision",
        "**Faster-RCNN Model** (confidence `%3.1f`) (overlap `%3.1f`)" % (confidence_threshold, overlap_threshold), 
        confidence_threshold, selected_classes)
    draw_legends(selected_classes)

# This sidebar UI is a little search engine to find certain object types.
def frame_selector_ui(summary):
    st.sidebar.markdown("# Frame")

    # The user can pick which type of object to search for.
    # object_type = st.sidebar.selectbox("Search for which objects?", summary.columns, 2)
    # The user can select a range for how many of the selected objecgt should be present.
    # min_elts, max_elts = st.sidebar.slider("How many %ss (select a range)?" % object_type, 0, 25, [10, 20])
    # selected_frames = get_selected_frames(summary, object_type, min_elts, max_elts)
    # len_selected_frames = len(selected_frames)
    len_selected_frames = len(IMG_NAMES)
    if len_selected_frames < 1:
        return None, None

    # Choose a frame out of the selected frames.
    selected_frame_index = st.sidebar.slider("Choose a frame (index)", 0, len_selected_frames - 1, 0)

    # Draw an altair chart in the sidebar with information on the frame.
    # objects_per_frame = summary.loc[selected_frames, object_type].reset_index(drop=True).reset_index()
    # chart = alt.Chart(objects_per_frame, height=120).mark_area().encode(
    #     alt.X("index:Q", scale=alt.Scale(nice=False)),
    #     alt.Y("%s:Q" % object_type))
    # selected_frame_df = pd.DataFrame({"selected_frame": [selected_frame_index]})
    # vline = alt.Chart(selected_frame_df).mark_rule(color="red").encode(x = "selected_frame")
    # st.sidebar.altair_chart(alt.layer(chart, vline))

    # selected_frame = selected_frames[selected_frame_index]
    selected_frame = IMG_NAMES[selected_frame_index]
    # ALL_LAYERS = {
    #     "t2": 0,
    #     "t10": 0,
    # }
    return selected_frame_index, selected_frame
def class_selector_ui():
    st.sidebar.markdown('## Select classes')
    selected_classes = [
        class_name for class_name in CLASS_NAMES
        if st.sidebar.checkbox(class_name, True)]
        # layer for layer_name, layer in ALL_LAYERS.items()
        # if st.sidebar.checkbox(layer_name, True)]
    return selected_classes

# Select frames based on the selection in the sidebar
@st.cache(hash_funcs={np.ufunc: str})
def get_selected_frames(summary, label, min_elts, max_elts):
    return summary[np.logical_and(summary[label] >= min_elts, summary[label] <= max_elts)].index

# This sidebar UI lets the user select parameters for the YOLO object detector.
def object_detector_ui():
    st.sidebar.markdown("# Model")
    confidence_threshold = st.sidebar.slider("Confidence threshold", 0.0, 1.0, 0.5, 0.01)
    overlap_threshold = st.sidebar.slider("Overlap threshold", 0.0, 1.0, 0.0, 0.01)
    return confidence_threshold, overlap_threshold

# Draws an image with boxes overlayed to indicate the presence of cars, pedestrians etc.
def draw_legends(selected_classes):
    for selected_class in selected_classes:
        cat_id = CLASS_NAMES.index(selected_class)
        color_bbox = PALETTES[cat_id]
        rec_size = 40
        rec = np.zeros([rec_size, rec_size, 3],dtype=np.uint8)
        rec[:,:] = tuple(color_bbox) 
        text = selected_class
        rec = cv2.copyMakeBorder( src=rec, top=0, bottom=0, left=0, right=int(rec_size*len(text)*1.1 ),borderType=0, value=[255]*3)
        cv2.putText(img=rec, text=text, org=(int(rec_size*1.1), int(rec_size*.8)), fontFace=2, fontScale=1, color=0, thickness=1, lineType=1)
        st.image(image=rec.astype(np.uint8))
        
def draw_image_with_boxes2(image, predict_dict, header, description, confidence_threshold, selected_classes):
    output = image.copy()
    score_list = predict_dict['score_list']
    bbox_list = predict_dict['bbox_list']
    pred_class_list = predict_dict['pred_class_list']
    pred_masks_list = predict_dict['pred_masks_list']
    pred_keypoints_list = predict_dict['pred_keypoints_list']
    vis_keypoints_list = predict_dict['vis_keypoints_list']
    kpt_confidences_list = predict_dict['kpt_confidences_list']
    for score, pred_class, bbox, mask, keypoints, vis_keypoints, kpt_confidences in zip(score_list,
                                                                                            pred_class_list,
                                                                                            bbox_list,
                                                                                            pred_masks_list,
                                                                                            pred_keypoints_list,
                                                                                            vis_keypoints_list,
                                                                                            kpt_confidences_list):
        if score > confidence_threshold:
            if pred_class in selected_classes:
                cat_id = CLASS_NAMES.index(pred_class)
                color_bbox = PALETTES[cat_id]
                output = draw_bbox(img=output, bbox=bbox, color=color_bbox,
                                show_label=False, text=None, thickness=5)
                # cv2.rectangle(img=image, pt1=(xmin, ymin), pt2=(
                # xmax, ymax), color=color_bbox, thickness=2)
    
    # Draw the header and image.
    st.subheader(header)
    st.markdown(description)
    st.image(image=output.astype(np.uint8),
            #  width=2000, 
             use_column_width=True
             )
    return output
def draw_image_with_boxes(image, boxes, header, description):
    # Superpose the semi-transparent object detection boxes.    # Colors for the boxes
    LABEL_COLORS = {
        "car": [255, 0, 0],
        "pedestrian": [0, 255, 0],
        "truck": [0, 0, 255],
        "trafficLight": [255, 255, 0],
        "biker": [255, 0, 255],
    }
    image_with_boxes = image.astype(np.float64)
    for _, (xmin, ymin, xmax, ymax, label) in boxes.iterrows():
        image_with_boxes[int(ymin):int(ymax),int(xmin):int(xmax),:] += LABEL_COLORS[label]
        image_with_boxes[int(ymin):int(ymax),int(xmin):int(xmax),:] /= 2

    # Draw the header and image.
    st.subheader(header)
    st.markdown(description)
    st.image(image_with_boxes.astype(np.uint8), use_column_width=True)

# Download a single file and make its content available as a string.
@st.cache(show_spinner=False)
def get_file_content_as_string(path):
    url = 'https://raw.githubusercontent.com/streamlit/demo-self-driving/master/' + path
    response = urllib.request.urlopen(url)
    return response.read().decode("utf-8")

# This function loads an image from Streamlit public repo on S3. We use st.cache on this
# function as well, so we can reuse the images across runs.
@st.cache(show_spinner=False)
def load_image(url):
    # with urllib.request.urlopen(url) as response:
    #     image = np.asarray(bytearray(response.read()), dtype="uint8")
    # image = cv2.imdecode(image, cv2.IMREAD_COLOR)
    # image = image[:, :, [2, 1, 0]] # BGR -> RGB
    image = cv2.imread(url)
    image = image[:, :, [2, 1, 0]] # BGR -> RGB
    return image


# Run the YOLO model to detect objects.
@st.cache(show_spinner=False)
def rcnn(image, confidence_threshold, overlap_threshold):
    weights_path = '/home/jitesh/3d/data/coco_data/ed/t2-5_coco-data/weights/faster_50_s2048_5/model_0001999.pth' 
    weights_path = '/home/jitesh/3d/data/coco_data/ed/t2-12_coco-data/weights/faster_50_s1024_3/model_0002999.pth' 
    weights_path = '/home/jitesh/3d/data/coco_data/ed/k/130_prediction_and_model_weight/weight/model_0001799.pth' 
    size = 1024
    inferer = D2Inferer(
        weights_path=weights_path,
        confidence_threshold=confidence_threshold,
        # class_names=['t2','t3','t4','t8'],
        class_names=CLASS_NAMES,
        # model='COCO-Detection/faster_rcnn_R_50_FPN_1x',
        model='Misc/cascade_mask_rcnn_R_50_FPN_3x',
        size_min=2340, size_max=3500,
        detectron2_dir_path="/home/jitesh/prj/detectron2",
        gray_on=True,
        # crop_mode=3,
        # crop_mode3_sizes=[1024],
        # crop_mode3_overlaps=[100]
    )
    inferer.cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = overlap_threshold
    predict_dict = inferer.predict(img=image)
    # score_list = predict_dict['score_list']
    # bbox_list = predict_dict['bbox_list']
    # pred_class_list = predict_dict['pred_class_list']
    # pred_masks_list = predict_dict['pred_masks_list']
    # pred_keypoints_list = predict_dict['pred_keypoints_list']
    # vis_keypoints_list = predict_dict['vis_keypoints_list']
    # kpt_confidences_list = predict_dict['kpt_confidences_list']
    # boxes = pd.DataFrame()
    # boxes = pd.DataFrame({"xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax, "labels": labels})
    return predict_dict
# Run the YOLO model to detect objects.
def yolo_v3(image, confidence_threshold, overlap_threshold):
    # Load the network. Because this is cached it will only happen once.
    @st.cache(allow_output_mutation=True)
    def load_network(config_path, weights_path):
        net = cv2.dnn.readNetFromDarknet(config_path, weights_path)
        output_layer_names = net.getLayerNames()
        output_layer_names = [output_layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
        return net, output_layer_names
    net, output_layer_names = load_network("yolov3.cfg", "yolov3.weights")

    # Run the YOLO neural net.
    blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416), swapRB=True, crop=False)
    net.setInput(blob)
    layer_outputs = net.forward(output_layer_names)

    # Supress detections in case of too low confidence or too much overlap.
    boxes, confidences, class_IDs = [], [], []
    H, W = image.shape[:2]
    for output in layer_outputs:
        for detection in output:
            scores = detection[5:]
            classID = np.argmax(scores)
            confidence = scores[classID]
            if confidence > confidence_threshold:
                box = detection[0:4] * np.array([W, H, W, H])
                centerX, centerY, width, height = box.astype("int")
                x, y = int(centerX - (width / 2)), int(centerY - (height / 2))
                boxes.append([x, y, int(width), int(height)])
                confidences.append(float(confidence))
                class_IDs.append(classID)
    indices = cv2.dnn.NMSBoxes(boxes, confidences, confidence_threshold, overlap_threshold)

    # Map from YOLO labels to Udacity labels.
    UDACITY_LABELS = {
        0: 'pedestrian',
        1: 'biker',
        2: 'car',
        3: 'biker',
        5: 'truck',
        7: 'truck',
        9: 'trafficLight'
    }
    xmin, xmax, ymin, ymax, labels = [], [], [], [], []
    if len(indices) > 0:
        # loop over the indexes we are keeping
        for i in indices.flatten():
            label = UDACITY_LABELS.get(class_IDs[i], None)
            if label is None:
                continue

            # extract the bounding box coordinates
            x, y, w, h = boxes[i][0], boxes[i][1], boxes[i][2], boxes[i][3]

            xmin.append(x)
            ymin.append(y)
            xmax.append(x+w)
            ymax.append(y+h)
            labels.append(label)

    boxes = pd.DataFrame({"xmin": xmin, "ymin": ymin, "xmax": xmax, "ymax": ymax, "labels": labels})
    return boxes[["xmin", "ymin", "xmax", "ymax", "labels"]]

# Path to the Streamlit public S3 bucket
DATA_URL_ROOT = "https://streamlit-self-driving.s3-us-west-2.amazonaws.com/"

# External files to download.
EXTERNAL_DEPENDENCIES = {
    # "yolov3.weights": {
    #     "url": "https://pjreddie.com/media/files/yolov3.weights",
    #     "size": 248007048
    # },
    # "yolov3.cfg": {
    #     "url": "https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg",
    #     "size": 8342
    # }
}

IMG_DATA_PATH = "data/annotated_data"
IMG_EXTENTIONS = ['jpg','jpeg', 'bmp', 'png', 'gif']
IMG_NAMES = sorted([fn for fn in os.listdir(IMG_DATA_PATH)
                if any(fn.endswith(ext) for ext in IMG_EXTENTIONS)])
CLASS_NAMES = ['t2', 't10']
CLASS_NAMES = ['bracket', 'distribution_board', 'distributor', 'downlight', 'interphone', 'light', 'outlet', 'plumbing', 'spotlight', 'wiring']
PALETTES = np.array(color_palette(
            # palette='Set2', 
            palette='hls', 
            n_colors=len(CLASS_NAMES)+1))*255
if __name__ == "__main__":
    main()

