import pydicom
from pydicom.dataset import Dataset
from pydicom.dataset import FileMetaDataset
from pydicom.uid import MediaStorageDirectoryStorage, EncapsulatedPDFStorage, generate_uid
import matplotlib
matplotlib.use("Agg")  # Use non-GUI backend to avoid Tkinter issues
import matplotlib.pyplot as plt  # Now import pyplot
from reportlab.pdfgen import canvas
from datetime import datetime, date


def extract_measurements(sr):
    """Extracts measurement annotations from an SR."""
    measurements = []
    probabilities = []

    if "ContentSequence" in sr:
        for itemLevel1 in sr.ContentSequence:
            if len(itemLevel1.ConceptNameCodeSequence) ==  1:
                if itemLevel1.ConceptNameCodeSequence[0].CodeMeaning == "Imaging Measurements":
                    for itemLevel2 in itemLevel1.ContentSequence:
                        for itemLevel3 in itemLevel2.ContentSequence:
                            if itemLevel3.ValueType == "SCOORD":
                                measurements.append(itemLevel3.GraphicData)
                            elif itemLevel3.ValueType == "NUM":
                                if len(itemLevel3.MeasuredValueSequence) == 1:
                                    probabilities.append(itemLevel3.MeasuredValueSequence[0].NumericValue)
    return measurements, probabilities


def overlay_measurements(image, measurements, probabilities):
    """Overlays extracted measurements onto the mammography image."""
    fig, ax = plt.subplots()
    ax.imshow(image, cmap='gray')

    # Draw each polyline
    for i in range(0, len(measurements), 1):
        measurement = measurements[i]
        x = measurement[0::2]  # Extract x-coordinates (every other value)
        y = measurement[1::2]  # Extract y-coordinates (every other value)
        ax.plot(x, y, 'lime', linewidth=1)  # Plot the entire polyline at once
        ax.text(x[-3] + 100, y[-3], f"{probabilities[i]:.2f} %", color='lime', fontsize=8)

    ax.axis("off")

    # Save the overlay as an image
    plt.savefig("temp.png", bbox_inches='tight', pad_inches=0)
    plt.close(fig)


def create_pdf(temp_image_path, measurements, sr, pdf_path):
    """Creates a PDF with the mammography image and extracted measurements."""
    c = canvas.Canvas(pdf_path)

    # Set font for the title
    c.setFont("Helvetica-Bold", 16)

    # Get page width to center the title
    page_width = 595  # Default A4 width in points
    title = "Mammography Report"
    c.drawCentredString(page_width / 2, 820, title)  # Adjust Y-position as needed

    # Reset font for other text
    c.setFont("Helvetica", 12)

    # Add patient info to the PDF
    c.drawString(70, 800, f"Patient ID: {sr.PatientID}")
    c.drawString(70, 785, f"Patient name: {sr.PatientName}")
    c.drawString(70, 770, f"Patient birth date: {formateted_datetime(sr.PatientBirthDate)}")
    c.drawString(70, 755, f"Patient sex: {sr.PatientSex}")
    c.drawString(70, 730, f"Study date: {formateted_datetime(sr.StudyDate, sr.StudyTime)}")
    c.drawString(70, 715, f"Report date: {formateted_datetime(sr.SeriesDate, sr.SeriesTime)}")
    c.drawString(70, 700, f"Referring physician: {sr.ReferringPhysicianName}")

    # Add the image to the PDF
    c.drawImage(temp_image_path, 70, 300)

    c.save()

# Convert DICOM date
def formateted_datetime(dicom_date, dicom_time = None):

    if dicom_date is None or dicom_date == '':
        return ''

    # Convert DICOM date
    formatted_date = datetime.strptime(dicom_date, "%Y%m%d").strftime("%Y-%m-%d")

    if dicom_time is None or dicom_time == '':
        return formatted_date

    # Convert DICOM time (handling optional fractions of a second)
    if "." in dicom_time:
        formatted_time = datetime.strptime(dicom_time, "%H%M%S.%f").strftime("%H:%M:%S.%f")[:-3]  # Keep milliseconds
    else:
        formatted_time = datetime.strptime(dicom_time, "%H%M%S").strftime("%H:%M:%S")

    # Combined datetime
    return f"{formatted_date} {formatted_time}"

def create_dcm_pdf(sr, pdf_path):
    ds = Dataset()

    # Add general DICOM metadata
    ds.PatientName = sr.PatientName
    ds.PatientID = sr.PatientID
    ds.PatientBirthDate = sr.PatientBirthDate
    ds.PatientSex = sr.PatientSex

    ds.StudyInstanceUID = sr.StudyInstanceUID
    ds.StudyDate = sr.StudyDate
    ds.StudyTime = sr.StudyTime
    ds.AccessionNumber = sr.AccessionNumber
    ds.ReferringPhysicianName = sr.ReferringPhysicianName
    ds.StudyID = sr.StudyID

    ds.SeriesInstanceUID = generate_uid()
    ds.SeriesDate = sr.SeriesDate
    ds.SeriesTime = sr.SeriesTime
    ds.SeriesNumber = 1
    ds.Modality = "DOC"

    ds.Manufacturer = "MammographyAI"
    ds.ConversionType = "DI"

    ds.SOPInstanceUID = generate_uid()
    ds.SOPClassUID = EncapsulatedPDFStorage

    # Open the PDF file and read it as binary data
    with open(pdf_path, 'rb') as f:
        pdf_data = f.read()

    # Add the EncapsulatedDocument (PDF content) to the DICOM dataset
    ds.ContentDate = ds.SeriesDate
    ds.ContentTime = ds.SeriesTime
    ds.AcquisitionDateTime = ""
    ds.InstanceNumber = 1
    ds.BurnedInAnnotation = "YES"
    ds.DocumentTitle = ""
    ds.EncapsulatedDocument = pdf_data
    ds.MIMETypeOfEncapsulatedDocument = "application/pdf"

    # Create a FileMetaDataset for DICOM file meta information
    file_meta = FileMetaDataset()
    file_meta.MediaStorageSOPClassUID = EncapsulatedPDFStorage
    file_meta.MediaStorageSOPInstanceUID = ds.SOPInstanceUID
    file_meta.TransferSyntaxUID = pydicom.uid.ImplicitVRLittleEndian
    file_meta.FileMetaInformationGroupLength = 0

    # Assign the file meta information to the dataset
    ds.file_meta = file_meta

    # Ensure preamble and "DICM" prefix is included
    ds.is_implicit_VR = True  # Set to explicit VR
    ds.is_little_endian = True  # Set to little endian
    
    return ds

def create(image, sr):
    measurements, probabilities = extract_measurements(sr)
    overlay_measurements(image, measurements, probabilities)
    create_pdf("temp.png", measurements, sr, "temp.pdf")
    return create_dcm_pdf(sr, "temp.pdf",)