#==============================================================================#
#  Author:       Dominik Müller                                                #
#  Copyright:    2020 IT-Infrastructure for Translational Medical Research,    #
#                University of Augsburg                                        #
#                                                                              #
#  This program is free software: you can redistribute it and/or modify        #
#  it under the terms of the GNU General Public License as published by        #
#  the Free Software Foundation, either version 3 of the License, or           #
#  (at your option) any later version.                                         #
#                                                                              #
#  This program is distributed in the hope that it will be useful,             #
#  but WITHOUT ANY WARRANTY; without even the implied warranty of              #
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the               #
#  GNU General Public License for more details.                                #
#                                                                              #
#  You should have received a copy of the GNU General Public License           #
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.       #
#==============================================================================#
#-----------------------------------------------------#
#                   Library imports                   #
#-----------------------------------------------------#
import requests
from tqdm import tqdm
import os
import zipfile

#-----------------------------------------------------#
#                    Configurations                   #
#-----------------------------------------------------#
# Data directory
path_data = "data"
# Links to the data set
url_vol = "https://zenodo.org/record/3757476/files/COVID-19-CT-Seg_20cases.zip?download=1"
url_seg = "https://zenodo.org/record/3757476/files/Lung_and_Infection_Mask.zip?download=1"

#-----------------------------------------------------#
#                  Download Function                  #
#-----------------------------------------------------#
# Author: Shenghan Gao (wy193777)
# Modifications: MCrazy
# Source: https://gist.github.com/wy193777/0e2a4932e81afc6aa4c8f7a2984f34e2
def download_from_url(url, dst):
    """
    @param: url to download file
    @param: dst place to put the file
    """
    file_size = int(requests.head(url).headers["Content-Length"])
    if os.path.exists(dst):
        first_byte = os.path.getsize(dst)
    else:
        first_byte = 0
    if first_byte >= file_size:
        print("WARNING: Skipping download due to files are already there.")
        return file_size
    header = {"Range": "bytes=%s-%s" % (first_byte, file_size)}
    pbar = tqdm(
        total=file_size, initial=first_byte,
        unit='B', unit_scale=True, desc=url.split('/')[-1])
    req = requests.get(url, headers=header, stream=True)
    with(open(dst, 'ab')) as f:
        for chunk in req.iter_content(chunk_size=1024):
            if chunk:
                f.write(chunk)
                pbar.update(1024)
    pbar.close()
    return file_size

#-----------------------------------------------------#
#                     Runner Code                     #
#-----------------------------------------------------#
# Create data structure
if not os.path.exists(path_data) : os.makedirs(path_data)

# Download CT volumes and save them into the data directory
path_vol_zip = os.path.join(path_data, "volumes.zip")
print("INFO:", "Downloading Volumes")
download_from_url(url_vol, path_vol_zip)
# Download segmentations and save them into the data directory
path_seg_zip = os.path.join(path_data, "segmentations.zip")
print("INFO:", "Downloading Segmentations")
download_from_url(url_seg, path_seg_zip)

# Extract sample list from the ZIP file
print("INFO:", "Obtain sample list from the volumes ZIP file")
with zipfile.ZipFile(path_vol_zip, "r") as zip_vol:
    sample_list = zip_vol.namelist()

# Iterate over the sample list and extract each sample from the ZIP files
print("INFO:", "Extracting data from ZIP files")
for sample in tqdm(sample_list):
    # Skip if file does not end with nii.gz
    if not sample.endswith(".nii.gz") : continue
    # Create sample directory
    path_sample = os.path.join(path_data, sample[:-len(".nii.gz")])
    if not os.path.exists(path_sample) : os.makedirs(path_sample)
    # Extract volume and store file into the sample directory
    with zipfile.ZipFile(path_vol_zip, "r") as zip_vol:
        zip_vol.extract(sample, path_sample)
    os.rename(os.path.join(path_sample, sample),
              os.path.join(path_sample, "imaging.nii.gz"))
    # Extract segmentation and store file into the sample directory
    with zipfile.ZipFile(path_seg_zip, "r") as zip_seg:
        zip_seg.extract(sample, path_sample)
    os.rename(os.path.join(path_sample, sample),
              os.path.join(path_sample, "segmentation.nii.gz"))

# Remove ZIP files due to disk space reduction
os.remove(path_vol_zip)
os.remove(path_seg_zip)

# Final info to console
print("INFO:", "Finished file structure creation")