from dataclasses import dataclass
from pathlib import Path

import numpy as np
import requests
from nibabel import load, save, Nifti1Image
from tqdm import tqdm


@dataclass
class Item:
    url: str
    path: str
    severity: str


def _download(url: str, data: str, fname: str):
    resp = requests.post(url, data=data, stream=True)
    total = int(resp.headers.get('content-length', 0))
    with open(fname, 'wb') as file, tqdm(
            desc=fname,
            total=total,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
    ) as bar:
        for data in resp.iter_content(chunk_size=1024):
            size = file.write(data)
            bar.update(size)


def _download_and_convert(
        url: str, fname: Path, alias: str, name: str, institution: str,
        email: str, force_update: bool = False
):
    minc_fname = fname.with_suffix('.mnc.gz')
    fname.parent.mkdir(parents=True, exist_ok=True)
    if minc_fname.exists() and not force_update:
        print(
            f"Skipping {str(minc_fname)} download, "
            f"since it already exists."
        )
    else:
        _download(
            url=url,
            data=f'do_download_alias={alias}'
                 f'&format_value=minc'
                 f'&zip_value=gnuzip'
                 f'&who_name={name}'
                 f'&who_institution={institution}'
                 f'&who_email={email}'
                 f'&download_for_real=%5BStart+download%21%5D',
            fname=str(minc_fname)
        )
    nii_fname = fname.with_suffix('.nii.gz')
    if nii_fname.exists() and not force_update:
        print(
            f"Skipping conversion of {str(minc_fname)}, "
            f"since {str(nii_fname)} already exists."
        )
    else:
        try:
            minc = load(minc_fname)
            affine = np.array([[0, 0, 1, 0],
                               [0, 1, 0, 0],
                               [1, 0, 0, 0],
                               [0, 0, 0, 1]])
            out = Nifti1Image(minc.get_fdata(), affine=affine)
            save(out, nii_fname)
            print(f'Successfully converted to {str(nii_fname)}')
        except Exception as e:
            print(e)


def _download_data(base_dir, name, institution, email):
    modality = 'T2'
    slice_thickness = '1mm'
    noise_levels = ['pn0', 'pn1', 'pn3', 'pn5']
    intensity_non_uniformities = ['rf0', 'rf20', 'rf40']
    severities = [
        Item(
            url='https://brainweb.bic.mni.mcgill.ca/cgi/brainweb2',
            path='lesions/severe',
            severity='AI+msles2'
        ),
        Item(
            url='https://brainweb.bic.mni.mcgill.ca/cgi/brainweb1',
            path='normal',
            severity='ICBM+normal'
        )
    ]

    for item in severities:
        out_dir = (base_dir / item.path)
        for noise_level in noise_levels:
            for intensity_non_uniformity in intensity_non_uniformities:
                alias = f"{modality}+{item.severity}+{slice_thickness}+{noise_level}+{intensity_non_uniformity}"
                fname = out_dir / alias.replace("+", "_").lower()
                _download_and_convert(
                    url=item.url,
                    fname=fname,
                    alias=alias,
                    name=name,
                    institution=institution,
                    email=email
                )


def _download_labels(base_dir, name, institution, email):
    items = [
        Item(
            url="https://brainweb.bic.mni.mcgill.ca/cgi/brainweb1",
            path="normal",
            severity="phantom_1.0mm_normal_crisp"
        ),  # normal
        Item(
            url="https://brainweb.bic.mni.mcgill.ca/cgi/brainweb2",
            path="severe_lesions",
            severity="phantom_1.0mm_msles3_crisp"
        )  # severe
    ]

    gt_dir = (base_dir / 'groundtruth')
    for item in items:
        _download_and_convert(
            url=item.url,
            fname=gt_dir / item.path,
            alias=item.severity,
            name=name,
            institution=institution,
            email=email
        )


def download_brainweb_dataset(
        # default path for google colab example
        base_dir: Path = Path('/content/data/Brainweb'),
        name: str = "",
        institution: str = "",
        email: str = "",
):
    _download_data(base_dir, name, institution, email)
    _download_labels(base_dir, name, institution, email)