import argparse
import glob
import os

import numpy as np
import torch
from sklearn.metrics import accuracy_score

import models_torch as models
import utils


EXPERIMENT_DATA_DIR = "/tmp/mgr"


def inference(parameters, verbose=True) -> int:

    # resolve device
    device = torch.device(
        "cuda:{}".format(parameters["gpu_number"]) if parameters["device_type"] == "gpu"
        else "cpu"
    )

    # load input images
    datum_l_cc = utils.load_images(parameters['image_path'], 'L-CC')
    datum_r_cc = utils.load_images(parameters['image_path'], 'R-CC')
    datum_l_mlo = utils.load_images(parameters['image_path'], 'L-MLO')
    datum_r_mlo = utils.load_images(parameters['image_path'], 'R-MLO')

    # construct models and prepare data
    if parameters["model_type"] == 'cnn':
        model = models.BaselineBreastModel(device, nodropout_probability=1.0, gaussian_noise_std=0.0).to(device)
        model.load_state_dict(torch.load(parameters["model_path"]))
        x = {
            "L-CC": torch.Tensor(datum_l_cc).permute(0, 3, 1, 2).to(device),
            "L-MLO": torch.Tensor(datum_l_mlo).permute(0, 3, 1, 2).to(device),
            "R-CC": torch.Tensor(datum_r_cc).permute(0, 3, 1, 2).to(device),
            "R-MLO": torch.Tensor(datum_r_mlo).permute(0, 3, 1, 2).to(device),
        }
    elif parameters["model_type"] == 'histogram':
        model = models.BaselineHistogramModel(num_bins=parameters["bins_histogram"]).to(device)
        model.load_state_dict(torch.load(parameters["model_path"]))
        x = torch.Tensor(utils.histogram_features_generator([
            datum_l_cc, datum_r_cc, datum_l_mlo, datum_r_mlo
        ], parameters)).to(device)
    else:
        raise RuntimeError(parameters["model_type"])

    # run prediction
    with torch.no_grad():
        prediction_density = model(x).cpu().numpy()

    if verbose:
        # nicely prints out the predictions
        print('Density prediction:\n'
              '\tAlmost entirely fatty (0):\t\t\t' + str(prediction_density[0, 0]) + '\n'
              '\tScattered areas of fibroglandular density (1):\t' + str(prediction_density[0, 1]) + '\n'
              '\tHeterogeneously dense (2):\t\t\t' + str(prediction_density[0, 2]) + '\n'
              '\tExtremely dense (3):\t\t\t\t' + str(prediction_density[0, 3]) + '\n')

    return np.argmax(prediction_density[0])+1 # return density in scope 1 to 4


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Run Inference')
    parser.add_argument('model_type')
    parser.add_argument('--bins-histogram', default=50)
    parser.add_argument('--model-path', default=None)
    parser.add_argument('--device-type', default="cpu")
    # parser.add_argument('--image-path', default="images/")
    args = parser.parse_args()

    parameters_ = {
        "model_type": args.model_type,
        "bins_histogram": args.bins_histogram,
        "model_path": args.model_path,
        "device_type": args.device_type,
        # "image_path": args.image_path,
    }

    if parameters_["model_path"] is None:
        if args.model_type == "histogram":
            parameters_["model_path"] = "saved_models/BreastDensity_BaselineHistogramModel/model.p"
        if args.model_type == "cnn":
            parameters_["model_path"] = "saved_models/BreastDensity_BaselineBreastModel/model.p"

    predicted_values = []
    real_values = []

    predicted_values_two_classes = []
    real_values_two_classes = []

    two_classes_mapping = {1: 0, 2: 0, 3: 1, 4: 1}

    for dir in glob.glob(f"{EXPERIMENT_DATA_DIR}/*/"):
        parameters_["image_path"] = dir
        predicted_density = inference(parameters_)

        with open(os.path.join(dir, "density.txt")) as file:
            real_density = int(file.read())
        print(f"Predicted density: {predicted_density}")
        print(f"Real density: {real_density}\n")

        print(f"Predicted density (2 cls): {two_classes_mapping[predicted_density]}")
        print(f"Real density (2 cls): {two_classes_mapping[real_density]}\n")

        predicted_values.append(predicted_density)
        real_values.append(real_density)

        predicted_values_two_classes.append(two_classes_mapping[predicted_density])
        real_values_two_classes.append(two_classes_mapping[real_density])

    print(f"Total accuracy: {accuracy_score(real_values, predicted_values)}")
    print(f"Total accuracy two classes: {accuracy_score(real_values_two_classes, predicted_values_two_classes)}")


"""
python density_model_torch_custom.py histogram
python density_model_torch_custom.py cnn
"""

[((25, 17, 25, 68), 'utils.load_images', 'utils.load_images', ({(25, 35, 25, 59): "parameters['image_path']", (25, 61, 25, 67): '"""L-CC"""'}, {}), "(parameters['image_path'], 'L-CC')", False, 'import utils\n'), ((26, 17, 26, 68), 'utils.load_images', 'utils.load_images', ({(26, 35, 26, 59): "parameters['image_path']", (26, 61, 26, 67): '"""R-CC"""'}, {}), "(parameters['image_path'], 'R-CC')", False, 'import utils\n'), ((27, 18, 27, 70), 'utils.load_images', 'utils.load_images', ({(27, 36, 27, 60): "parameters['image_path']", (27, 62, 27, 69): '"""L-MLO"""'}, {}), "(parameters['image_path'], 'L-MLO')", False, 'import utils\n'), ((28, 18, 28, 70), 'utils.load_images', 'utils.load_images', ({(28, 36, 28, 60): "parameters['image_path']", (28, 62, 28, 69): '"""R-MLO"""'}, {}), "(parameters['image_path'], 'R-MLO')", False, 'import utils\n'), ((66, 13, 66, 65), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((96, 15, 96, 53), 'glob.glob', 'glob.glob', ({(96, 25, 96, 52): 'f"""{EXPERIMENT_DATA_DIR}/*/"""'}, {}), "(f'{EXPERIMENT_DATA_DIR}/*/')", False, 'import glob\n'), ((50, 9, 50, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((61, 11, 61, 43), 'numpy.argmax', 'np.argmax', ({(61, 21, 61, 42): 'prediction_density[0]'}, {}), '(prediction_density[0])', True, 'import numpy as np\n'), ((33, 30, 33, 66), 'torch.load', 'torch.load', ({(33, 41, 33, 65): "parameters['model_path']"}, {}), "(parameters['model_path'])", False, 'import torch\n'), ((32, 16, 32, 101), 'models_torch.BaselineBreastModel', 'models.BaselineBreastModel', (), '', True, 'import models_torch as models\n'), ((42, 30, 42, 66), 'torch.load', 'torch.load', ({(42, 41, 42, 65): "parameters['model_path']"}, {}), "(parameters['model_path'])", False, 'import torch\n'), ((100, 18, 100, 50), 'os.path.join', 'os.path.join', ({(100, 31, 100, 34): 'dir', (100, 36, 100, 49): '"""density.txt"""'}, {}), "(dir, 'density.txt')", False, 'import os\n'), ((114, 29, 114, 74), 'sklearn.metrics.accuracy_score', 'accuracy_score', ({(114, 44, 114, 55): 'real_values', (114, 57, 114, 73): 'predicted_values'}, {}), '(real_values, predicted_values)', False, 'from sklearn.metrics import accuracy_score\n'), ((115, 41, 115, 110), 'sklearn.metrics.accuracy_score', 'accuracy_score', ({(115, 56, 115, 79): 'real_values_two_classes', (115, 81, 115, 109): 'predicted_values_two_classes'}, {}), '(real_values_two_classes, predicted_values_two_classes)', False, 'from sklearn.metrics import accuracy_score\n'), ((41, 16, 41, 84), 'models_torch.BaselineHistogramModel', 'models.BaselineHistogramModel', (), '', True, 'import models_torch as models\n'), ((43, 25, 45, 22), 'utils.histogram_features_generator', 'utils.histogram_features_generator', ({(43, 60, 45, 9): '[datum_l_cc, datum_r_cc, datum_l_mlo, datum_r_mlo]', (45, 11, 45, 21): 'parameters'}, {}), '([datum_l_cc, datum_r_cc, datum_l_mlo,\n    datum_r_mlo], parameters)', False, 'import utils\n'), ((35, 20, 35, 44), 'torch.Tensor', 'torch.Tensor', ({(35, 33, 35, 43): 'datum_l_cc'}, {}), '(datum_l_cc)', False, 'import torch\n'), ((36, 21, 36, 46), 'torch.Tensor', 'torch.Tensor', ({(36, 34, 36, 45): 'datum_l_mlo'}, {}), '(datum_l_mlo)', False, 'import torch\n'), ((37, 20, 37, 44), 'torch.Tensor', 'torch.Tensor', ({(37, 33, 37, 43): 'datum_r_cc'}, {}), '(datum_r_cc)', False, 'import torch\n'), ((38, 21, 38, 46), 'torch.Tensor', 'torch.Tensor', ({(38, 34, 38, 45): 'datum_r_mlo'}, {}), '(datum_r_mlo)', False, 'import torch\n')]