# import pydicom

# import numpy as np

# from PIL import Image

#

# tolerance = 1

#

# # Load the DICOM file

# dicom = pydicom.dcmread(r"/home/ubuntu/nlp_project/Code/physionet.org/files/mimic-cxr/2.1.0/files/p10/p10000032/s50414267/02aa804e-bde0afdd-112c0b34-7bc16630-4e384014.dcm")

# dicom_pixels = dicom.pixel_array

#

# # Load the PNG file

# png_image = Image.open(r"/home/ubuntu/nlp_project/Code/physionet.org/files/mimic-cxr/2.1.0/files/p10/p10000032/s50414267/out_png/02aa804e-bde0afdd-112c0b34-7bc16630-4e384014.png")

# png_pixels = np.array(png_image)

#

# # Compare shapes

# assert dicom_pixels.shape == png_pixels.shape, "Shapes of DICOM and PNG do not match"

#

# # Compare pixel values

# difference = np.abs(dicom_pixels - png_pixels)

#

# dicom_non_zero_elements = dicom_pixels[dicom_pixels != 0]

# dicom_non_zero_count = len(dicom_non_zero_elements)

#

# # Get non-zero elements in PNG array

# png_non_zero_elements = png_pixels[png_pixels != 0]

# png_non_zero_count = len(png_non_zero_elements)

#

# print(f"Non-zero elements in DICOM: {dicom_non_zero_count}")

# print(f"Non-zero elements in PNG: {png_non_zero_count}")

#

# # If needed, print the actual non-zero elements (be cautious for large arrays)

# print(f"Non-zero elements in DICOM array:\n{dicom_non_zero_elements}")

# print(f"Non-zero elements in PNG array:\n{png_non_zero_elements}")

#

#

# print(f"First 10 non-zero elements in DICOM array:\n{dicom_non_zero_elements[:10]}")

# print(f"First 10 non-zero elements in PNG array:\n{png_non_zero_elements[:10]}")

#

#

# # assert np.all(difference < tolerance), f"Pixel difference exceeds tolerance: {difference.max()}"

import os

import pandas as pd

from tqdm import tqdm

import pydicom

import numpy as np

from PIL import Image

# Function to convert DICOM to PNG

def convert_dicom_to_png(dicom_path, output_path):

    try:

        # Read the DICOM file

        dicom = pydicom.dcmread(dicom_path)

        # Get pixel array

        pixel_array = dicom.pixel_array

        # Normalize pixel values to 0-255

        pixel_array = ((pixel_array - pixel_array.min()) / (pixel_array.max() - pixel_array.min()) * 255).astype(

            np.uint8)

        # Save as PNG

        image = Image.fromarray(pixel_array)

        image.save(output_path)

    except Exception as e:

        print(f"Error converting {dicom_path} to PNG: {e}")

# Function to extract findings and impressions from a report

def extract_findings_and_impression(file_path):

    with open(file_path, 'r') as file:

        content = file.read()

    # Extract Findings

    findings_start = content.find("FINDINGS:")

    impression_start = content.find("IMPRESSION:")

    findings = ""

    impression = ""

    if findings_start != -1:

        findings = content[findings_start + len("FINDINGS:"):impression_start].strip()

    if impression_start != -1:

        impression = content[impression_start + len("IMPRESSION:"):].strip()

    return findings, impression

# Main logic to create the DataFrame

reports_root_path = input("Enter the root path for reports: ").strip()

# Ensure the path exists

if not os.path.exists(reports_root_path):

    raise FileNotFoundError(f"The specified path does not exist: {reports_root_path}")

data = []

grp_folders = os.listdir(reports_root_path)

for p_grp in grp_folders:

    cxr_path = os.path.join(reports_root_path, p_grp)

    p_files = os.listdir(cxr_path)

    for p in p_files:

        res_path = os.path.join(cxr_path, p)

        if os.path.isdir(res_path):

            dicom_dirs = [d for d in os.listdir(res_path) if os.path.isdir(os.path.join(res_path, d))]

            txt_files = [f for f in os.listdir(res_path) if f.endswith('.txt') and f.startswith('s')]

            for dicom_dir in dicom_dirs:

                dicom_path = os.path.join(res_path, dicom_dir)

                dicom_files = [os.path.join(dicom_path, f) for f in os.listdir(dicom_path) if f.endswith('.dcm')]

                report_file = f"{dicom_dir}.txt"

                if report_file in txt_files:

                    report_path = os.path.join(res_path, report_file)

                    findings, impressions = extract_findings_and_impression(report_path)

                    for dicom_file in dicom_files:

                        dicom_id = os.path.basename(dicom_file)

                        png_path = dicom_file.replace('.dcm', '.png')  # Define the PNG output path

                        # Convert the DICOM to PNG

                        convert_dicom_to_png(dicom_file, png_path)

                        # Append data to the list

                        # data.append({

                        #     "dicom_path": dicom_file,

                        #     "png_path": png_path,

                        #     "dicom_id": dicom_id,

                        #     "findings": findings,

                        #     "impressions": impressions

                        # })

                        data_entry = {

                            "dicom_path": dicom_file,

                            "png_path": png_path,

                            "dicom_id": dicom_id,

                            "findings": findings,

                            "impressions": impressions

                        }

                        data.append(data_entry)

                        print(f"Processed PNG path: {data_entry['png_path']}")

df = pd.DataFrame(data)

print(df.head())

print(f"Total entries: {len(df)}")

# Save the DataFrame to a CSV file

df.to_csv('data_with_png_paths.csv', index=False)