import os

from django.db import models
import uuid

# Create your models here.
from djcelery_model.models import TaskMixin
from polymorphic.models import PolymorphicModel

from genui.utils.models import NON_POLYMORPHIC_CASCADE, OverwriteStorage
from genui.utils.extensions.tasks.models import TaskShortcutsMixIn, PolymorphicTaskManager
from genui.projects.models import DataSet


class AlgorithmMode(models.Model):
    name = models.CharField(unique=True, blank=False, max_length=32)

    def __str__(self):
        return '%s object (%s)' % (self.__class__.__name__, self.name)


class ModelFileFormat(models.Model):
    fileExtension = models.CharField(max_length=32, blank=False, unique=True)
    description = models.TextField(max_length=10000, blank=True)

class ImportableModelComponent(models.Model):

    corePackage = models.CharField(blank=False, null=False, default='genui.models.genuimodels', max_length=1024)

    class Meta:
        abstract = True


class Algorithm(ImportableModelComponent):
    name = models.CharField(blank=False, max_length=128, unique=True)
    fileFormats = models.ManyToManyField(ModelFileFormat)
    validModes = models.ManyToManyField(AlgorithmMode)

    def __str__(self):
        return '%s object (%s)' % (self.__class__.__name__, self.name)


class ModelParameter(models.Model):
    STRING = 'string'
    BOOL = 'bool'
    INTEGER = 'integer'
    FLOAT = 'float'
    CONTENT_TYPES = [
       (STRING, 'String'),
       (BOOL, 'Logical'),
       (INTEGER, 'Integer'),
       (FLOAT, 'Float'),
    ]

    name = models.CharField(max_length=128, blank=False)
    algorithm = models.ForeignKey(Algorithm, on_delete=models.CASCADE, null=False, related_name='parameters')
    contentType = models.CharField(max_length=32, choices=CONTENT_TYPES, default=STRING)
    defaultValue = models.ForeignKey("ModelParameterValue", on_delete=models.SET_NULL, null=True)

    class Meta:
        unique_together = ('name', 'algorithm')

    def __str__(self):
        return '%s object (%s)' % (self.__class__.__name__, self.name)

class ModelBuilder(ImportableModelComponent):
    name = models.CharField(max_length=128, blank=False, unique=True)

    def __str__(self):
        return '%s object (%s)' % (self.__class__.__name__, self.name)

class ModelFile(models.Model):
    MAIN = "main"
    AUXILIARY = "aux"
    KINDS = [
       (MAIN, 'Main'),
       (AUXILIARY, 'Auxiliary'),
    ]

    class Rejected(Exception):

        def __init__(self, msg):
            super().__init__(msg)

    class InvalidFileFormatError(Exception):

        def __init__(self, msg):
            super().__init__(msg)

    modelInstance = models.ForeignKey("Model", null=False, related_name="files", on_delete=models.CASCADE)
    kind = models.CharField(max_length=32, choices=KINDS, null=False, default=AUXILIARY)
    note = models.CharField(max_length=128, blank=True)
    format = models.ForeignKey(ModelFileFormat, null=True, on_delete=models.CASCADE)
    file = models.FileField(null=True, upload_to='models/', storage=OverwriteStorage()) # TODO: add custom logic to save in a directory specific to the project where the model is

    @property
    def path(self):
        return self.file.path

    @staticmethod
    def generateMainFileName(model, fileFormat):
        return f"{model.trainingStrategy.algorithm.name}{model.id}_project{model.project.id}_{uuid.uuid4().hex}_main{fileFormat.fileExtension}"

    @staticmethod
    def generateAuxFileName(model, fileFormat):
        return f"{model.trainingStrategy.algorithm.name}{model.id}_project{model.project.id}_{uuid.uuid4().hex}_aux{fileFormat.fileExtension}"

    @staticmethod
    def create(model, name, file_, kind=AUXILIARY, note=None):
        if not note:
            note = ''
        algorithm = model.trainingStrategy.algorithm
        if kind == ModelFile.MAIN and model.modelFile:
            file_format = None
            for format_ in algorithm.fileFormats.all():
                if name.endswith(format_.fileExtension):
                    file_format = format_
                    break
            if not file_format:
                raise ModelFile.InvalidFileFormatError(f"The extension for file '{name}' of the submitted file did not match any of the known formats for algorithm: ({algorithm.name}).")

            if model.modelFile.format.fileExtension == file_format.fileExtension:
                model.modelFile.file.save(os.path.basename(model.modelFile.path), file_)
            else:
                model.modelFile.delete()
                ModelFile.objects.create(
                    model=model,
                    kind=ModelFile.MAIN,
                    format=file_format,
                    note=note,
                    file=file_
                )

            return model.modelFile
        else:
            file_format = None
            for format_ in ModelFileFormat.objects.all():
                if name.endswith(format_.fileExtension):
                    file_format = format_
                    break
            if kind == ModelFile.MAIN:
                if not file_format:
                    raise ModelFile.InvalidFileFormatError(f"The extension for file '{name}' of the submitted file did not match any of the known formats for algorithm: ({algorithm.name}).")
                ret = ModelFile.objects.create(
                    modelInstance=model,
                    kind=ModelFile.MAIN,
                    format=file_format,
                    note=note
                )
                ret.file.save(ret.generateMainFileName(model, file_format), file_)
            else:
                ret = ModelFile.objects.create(
                    modelInstance=model,
                    kind=kind,
                    format=file_format if file_format else ModelFileFormat.objects.get_or_create(
                        fileExtension='.' + name.split('.')[-1]
                    )[0],
                    note=note
                )
                ret.file.save(ret.generateAuxFileName(model, ret.format), file_)

            return ret

class Model(TaskShortcutsMixIn, TaskMixin, DataSet):
    objects = PolymorphicTaskManager()
    builder = models.ForeignKey(ModelBuilder, on_delete=models.CASCADE, null=False)

    def __str__(self):
        return '%s object (%s)' % (self.__class__.__name__, self.name)

    @property
    def modelFile(self):
        # TODO: exception when more than one main file found
        main = self.files.filter(kind=ModelFile.MAIN)
        if main:
            return main.get()
        else:
            return None

    def onFileSave(self, saved : ModelFile):
        """
        This will be called when a file is being
        saved to this model instance. You can throw
        the ModelFile.Rejected exception if the file
        is invalid.

        :param saved:
        :return:
        """

        pass

    # @modelFile.setter
    # def modelFile(self, val):
    #     main = self.files.filter(kind=ModelFile.MAIN)
    #     if main:
    #         main.delete()
    #     val.kind = ModelFile.MAIN
    #     val.save()
    #     self.files.add(val)
    #     self.save()

    @property
    def trainingStrategy(self):
        count = self.trainingStrategies.count()
        if count == 1:
            return self.trainingStrategies.get()
        elif count == 0:
            return None
        else:
            raise Exception("Training strategy returned more than one value. This indicates an integrity error in the database!")

    @property
    def validationStrategy(self):
        count = self.validationStrategies.count()
        if count == 1:
            return self.validationStrategies.get()
        elif count == 0:
            return None
        else:
            raise Exception("Validation strategy returned more than one value. This indicates an integrity error in the database!")


class TrainingStrategy(PolymorphicModel):
    algorithm = models.ForeignKey(Algorithm, on_delete=models.CASCADE, null=False)
    mode = models.ForeignKey(AlgorithmMode, on_delete=models.CASCADE, null=False)
    modelInstance = models.ForeignKey(Model, null=False, on_delete=models.CASCADE, related_name="trainingStrategies")


class ModelParameterValue(PolymorphicModel):
    parameter = models.ForeignKey(ModelParameter, on_delete=models.CASCADE, null=False)
    strategy = models.ForeignKey(TrainingStrategy, on_delete=NON_POLYMORPHIC_CASCADE, null=True, related_name='parameters')

    @staticmethod
    def parseValue(val):
        return str(val)


class ModelParameterStr(ModelParameterValue):
    value = models.CharField(max_length=1024)


class ModelParameterBool(ModelParameterValue):
    value = models.BooleanField(null=False)

    @staticmethod
    def parseValue(val):
        return bool(val)


class ModelParameterInt(ModelParameterValue):
    value = models.IntegerField(null=False)

    @staticmethod
    def parseValue(val):
        return int(val)


class ModelParameterFloat(ModelParameterValue):
    value = models.FloatField(null=False)

    @staticmethod
    def parseValue(val):
        return float(val)

PARAM_VALUE_CTYPE_TO_MODEL_MAP = {
    ModelParameter.STRING : ModelParameterStr,
    ModelParameter.INTEGER : ModelParameterInt,
    ModelParameter.FLOAT : ModelParameterFloat,
    ModelParameter.BOOL : ModelParameterBool
}


class ModelPerformanceMetric(ImportableModelComponent):
    name = models.CharField(unique=True, blank=False, max_length=128)
    validModes = models.ManyToManyField(AlgorithmMode, related_name='metrics')
    validAlgorithms = models.ManyToManyField(Algorithm, related_name='metrics')
    description = models.TextField(max_length=10000, blank=True)

    def __str__(self):
        return '%s object (%s)' % (self.__class__.__name__, self.name)


class ValidationStrategy(PolymorphicModel):
    metrics = models.ManyToManyField(ModelPerformanceMetric)
    modelInstance = models.ForeignKey(Model, null=False, on_delete=models.CASCADE, related_name='validationStrategies')


class CV(ValidationStrategy):
    cvFolds = models.IntegerField(blank=False)

    class Meta:
        abstract = True


class ValidationSet(ValidationStrategy):
    validSetSize = models.FloatField(blank=False)

    class Meta:
        abstract = True


class BasicValidationStrategy(ValidationSet, CV):
    pass


class ModelPerformance(PolymorphicModel):
    metric = models.ForeignKey(ModelPerformanceMetric, null=False, on_delete=models.CASCADE)
    value = models.FloatField(blank=False)
    model = models.ForeignKey(Model, null=False, on_delete=NON_POLYMORPHIC_CASCADE, related_name="performance")


class ModelPerformanceCV(ModelPerformance):
    fold = models.IntegerField(blank=False)

class ModelPerfomanceNN(ModelPerformance):
    epoch = models.IntegerField(null=False, blank=False)
    step = models.IntegerField(null=False, blank=False)

class ROCCurvePoint(ModelPerformance):

    fpr = models.FloatField(blank=False)
    auc = models.ForeignKey(ModelPerformance, null=False, on_delete=NON_POLYMORPHIC_CASCADE, related_name="points")

    @property
    def tpr(self):
        return self.value
[((16, 11, 16, 68), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((23, 20, 23, 77), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((24, 18, 24, 64), 'django.db.models.TextField', 'models.TextField', (), '', False, 'from django.db import models\n'), ((28, 18, 28, 112), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((35, 11, 35, 69), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((36, 18, 36, 57), 'django.db.models.ManyToManyField', 'models.ManyToManyField', ({(36, 41, 36, 56): 'ModelFileFormat'}, {}), '(ModelFileFormat)', False, 'from django.db import models\n'), ((37, 17, 37, 54), 'django.db.models.ManyToManyField', 'models.ManyToManyField', ({(37, 40, 37, 53): 'AlgorithmMode'}, {}), '(AlgorithmMode)', False, 'from django.db import models\n'), ((55, 11, 55, 56), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((56, 16, 56, 109), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((57, 18, 57, 88), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((58, 19, 58, 97), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((67, 11, 67, 69), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((90, 20, 90, 106), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((91, 11, 91, 88), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((92, 11, 92, 55), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((93, 13, 93, 84), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((165, 14, 165, 38), 'genui.utils.extensions.tasks.models.PolymorphicTaskManager', 'PolymorphicTaskManager', ({}, {}), '()', False, 'from genui.utils.extensions.tasks.models import TaskShortcutsMixIn, PolymorphicTaskManager\n'), ((166, 14, 166, 83), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((225, 16, 225, 82), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((226, 11, 226, 81), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((227, 20, 227, 117), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((231, 16, 231, 87), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((232, 15, 232, 123), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((240, 12, 240, 45), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((244, 12, 244, 43), 'django.db.models.BooleanField', 'models.BooleanField', (), '', False, 'from django.db import models\n'), ((252, 12, 252, 43), 'django.db.models.IntegerField', 'models.IntegerField', (), '', False, 'from django.db import models\n'), ((260, 12, 260, 41), 'django.db.models.FloatField', 'models.FloatField', (), '', False, 'from django.db import models\n'), ((275, 11, 275, 69), 'django.db.models.CharField', 'models.CharField', (), '', False, 'from django.db import models\n'), ((276, 17, 276, 78), 'django.db.models.ManyToManyField', 'models.ManyToManyField', (), '', False, 'from django.db import models\n'), ((277, 22, 277, 79), 'django.db.models.ManyToManyField', 'models.ManyToManyField', (), '', False, 'from django.db import models\n'), ((278, 18, 278, 64), 'django.db.models.TextField', 'models.TextField', (), '', False, 'from django.db import models\n'), ((285, 14, 285, 60), 'django.db.models.ManyToManyField', 'models.ManyToManyField', ({(285, 37, 285, 59): 'ModelPerformanceMetric'}, {}), '(ModelPerformanceMetric)', False, 'from django.db import models\n'), ((286, 20, 286, 119), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((290, 14, 290, 46), 'django.db.models.IntegerField', 'models.IntegerField', (), '', False, 'from django.db import models\n'), ((297, 19, 297, 49), 'django.db.models.FloatField', 'models.FloatField', (), '', False, 'from django.db import models\n'), ((308, 13, 308, 92), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((309, 12, 309, 42), 'django.db.models.FloatField', 'models.FloatField', (), '', False, 'from django.db import models\n'), ((310, 12, 310, 111), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((314, 11, 314, 43), 'django.db.models.IntegerField', 'models.IntegerField', (), '', False, 'from django.db import models\n'), ((317, 12, 317, 56), 'django.db.models.IntegerField', 'models.IntegerField', (), '', False, 'from django.db import models\n'), ((318, 11, 318, 55), 'django.db.models.IntegerField', 'models.IntegerField', (), '', False, 'from django.db import models\n'), ((322, 10, 322, 40), 'django.db.models.FloatField', 'models.FloatField', (), '', False, 'from django.db import models\n'), ((323, 10, 323, 115), 'django.db.models.ForeignKey', 'models.ForeignKey', (), '', False, 'from django.db import models\n'), ((94, 68, 94, 86), 'genui.utils.models.OverwriteStorage', 'OverwriteStorage', ({}, {}), '()', False, 'from genui.utils.models import NON_POLYMORPHIC_CASCADE, OverwriteStorage\n'), ((102, 94, 102, 106), 'uuid.uuid4', 'uuid.uuid4', ({}, {}), '()', False, 'import uuid\n'), ((106, 94, 106, 106), 'uuid.uuid4', 'uuid.uuid4', ({}, {}), '()', False, 'import uuid\n'), ((123, 42, 123, 80), 'os.path.basename', 'os.path.basename', ({(123, 59, 123, 79): 'model.modelFile.path'}, {}), '(model.modelFile.path)', False, 'import os\n')]