from typing import NamedTuple
import os
import torch
import numpy as np
import pydicom
from service.structs import PredictorInput, TagError, MetaTags


class InputImage(NamedTuple):
    img: np.ndarray
    ww: int
    wc: int
    uid: str
    color_inversion: bool


def _get_ww_wc(study_ww, study_wc):
    if type(study_ww) is pydicom.valuerep.DSfloat:
        return int(study_ww), int(study_wc)
    return int(study_ww[0]), int(study_wc[0])


def _is_color_inverted(study):
    return study.PhotometricInterpretation == "MONOCHROME1"


def _prep_img(img_pack: InputImage) -> torch.Tensor:
    img = img_pack.img.astype(np.float32)
    lower_bound = img_pack.wc - 0.5 * img_pack.ww
    upper_bound = img_pack.wc + 0.5 * img_pack.ww
    img = np.clip(img, lower_bound, upper_bound)
    img = ((img - lower_bound) / img_pack.ww) * 255
    img = img.astype(np.uint8)
    if img_pack.color_inversion:
        img = 255 - img
    return torch.from_numpy(img[None, ...])


def _check_required_tags(study: pydicom.FileDataset, pathology: str):
    invalid_tags = []
    if not (hasattr(study, "PhotometricInterpretation") and 
            study.PhotometricInterpretation in ("MONOCHROME1", "MONOCHROME2")):
        invalid_tags.append("PhotometricInterpretation")
    if pathology in ("shoulder", "wrist"):
        if not (hasattr(study, "PixelSpacing") or
                hasattr(study, "ImagerPixelSpacing")):
            invalid_tags.append("PixelSpacing/ImagerPixelSpacing")
    for tag in ("WindowWidth", "WindowCenter"):
        if not hasattr(study, tag) or getattr(study, tag) == "":
            invalid_tags.append(tag)
    if len(invalid_tags) == 1:
        raise TagError(f"DICOM тег {invalid_tags[0]} не заполнен, "\
                       "либо имеет некорректное значение")
    elif len(invalid_tags) > 1:
        raise TagError(f"DICOM теги {", ".join(invalid_tags)} не заполнены, "\
                       "либо имеют некорректные значения")


def _get_meta_tags(study: pydicom.FileDataset):
    return MetaTags(study.StudyInstanceUID, study.SeriesInstanceUID,
                    getattr(study, "PatientID", None), 
                    getattr(study, "AccessionNumber", None),
                    getattr(study, "IssuerOfPatientID", None),
                    getattr(study, "FillerOrderNumberImagingServiceRequest",
                                    None))


def _get_px_size(study: pydicom.FileDataset):
    if hasattr(study, "PixelSpacing"):
        return study.PixelSpacing
    if hasattr(study, "ImagerPixelSpacing"):
        return study.ImagerPixelSpacing
    return None, None


def prep_imgs(pathology: str, study_path: str) -> tuple[MetaTags, PredictorInput]:
    instance = pydicom.dcmread(study_path, force=True)
    _check_required_tags(instance, pathology)
    meta_tags = _get_meta_tags(instance)
    ww, wc = _get_ww_wc(instance.WindowWidth, instance.WindowCenter)
    clr_inverted = _is_color_inverted(instance)
    img = InputImage(instance.pixel_array, ww, wc, meta_tags.study_iuid, clr_inverted)
    pred_input = PredictorInput(meta_tags.study_iuid, _prep_img(img),
                                getattr(instance, "ImageLaterality", None),
                                *_get_px_size(instance))
    return meta_tags, pred_input