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
[((427, 484), 'django.db.models.CharField', 'models.CharField', ([], {'unique': '(True)', 'blank': '(False)', 'max_length': '(32)'}), '(unique=True, blank=False, max_length=32)\n', (443, 484), False, 'from django.db import models\n'), ((639, 696), 'django.db.models.CharField', 'models.CharField', ([], {'max_length': '(32)', 'blank': '(False)', 'unique': '(True)'}), '(max_length=32, blank=False, unique=True)\n', (655, 696), False, 'from django.db import models\n'), ((715, 761), 'django.db.models.TextField', 'models.TextField', ([], {'max_length': '(10000)', 'blank': '(True)'}), '(max_length=10000, blank=True)\n', (731, 761), False, 'from django.db import models\n'), ((828, 927), 'django.db.models.CharField', 'models.CharField', ([], {'blank': '(False)', 'null': '(False)', 'default': '"""genui.models.genuimodels"""', 'max_length': '(1024)'}), "(blank=False, null=False, default=\n    'genui.models.genuimodels', max_length=1024)\n", (844, 927), False, 'from django.db import models\n'), ((1020, 1078), 'django.db.models.CharField', 'models.CharField', ([], {'blank': '(False)', 'max_length': '(128)', 'unique': '(True)'}), '(blank=False, max_length=128, unique=True)\n', (1036, 1078), False, 'from django.db import models\n'), ((1097, 1136), 'django.db.models.ManyToManyField', 'models.ManyToManyField', (['ModelFileFormat'], {}), '(ModelFileFormat)\n', (1119, 1136), False, 'from django.db import models\n'), ((1154, 1191), 'django.db.models.ManyToManyField', 'models.ManyToManyField', (['AlgorithmMode'], {}), '(AlgorithmMode)\n', (1176, 1191), False, 'from django.db import models\n'), ((1556, 1601), 'django.db.models.CharField', 'models.CharField', ([], {'max_length': '(128)', 'blank': '(False)'}), '(max_length=128, blank=False)\n', (1572, 1601), False, 'from django.db import models\n'), ((1618, 1715), 'django.db.models.ForeignKey', 'models.ForeignKey', (['Algorithm'], {'on_delete': 'models.CASCADE', 'null': '(False)', 'related_name': '"""parameters"""'}), "(Algorithm, on_delete=models.CASCADE, null=False,\n    related_name='parameters')\n", (1635, 1715), False, 'from django.db import models\n'), ((1730, 1800), 'django.db.models.CharField', 'models.CharField', ([], {'max_length': '(32)', 'choices': 'CONTENT_TYPES', 'default': 'STRING'}), '(max_length=32, choices=CONTENT_TYPES, default=STRING)\n', (1746, 1800), False, 'from django.db import models\n'), ((1820, 1898), 'django.db.models.ForeignKey', 'models.ForeignKey', (['"""ModelParameterValue"""'], {'on_delete': 'models.SET_NULL', 'null': '(True)'}), "('ModelParameterValue', on_delete=models.SET_NULL, null=True)\n", (1837, 1898), False, 'from django.db import models\n'), ((2117, 2175), 'django.db.models.CharField', 'models.CharField', ([], {'max_length': '(128)', 'blank': '(False)', 'unique': '(True)'}), '(max_length=128, blank=False, unique=True)\n', (2133, 2175), False, 'from django.db import models\n'), ((2654, 2745), 'django.db.models.ForeignKey', 'models.ForeignKey', (['"""Model"""'], {'null': '(False)', 'related_name': '"""files"""', 'on_delete': 'models.CASCADE'}), "('Model', null=False, related_name='files', on_delete=\n    models.CASCADE)\n", (2671, 2745), False, 'from django.db import models\n'), ((2752, 2829), 'django.db.models.CharField', 'models.CharField', ([], {'max_length': '(32)', 'choices': 'KINDS', 'null': '(False)', 'default': 'AUXILIARY'}), '(max_length=32, choices=KINDS, null=False, default=AUXILIARY)\n', (2768, 2829), False, 'from django.db import models\n'), ((2841, 2885), 'django.db.models.CharField', 'models.CharField', ([], {'max_length': '(128)', 'blank': '(True)'}), '(max_length=128, blank=True)\n', (2857, 2885), False, 'from django.db import models\n'), ((2899, 2970), 'django.db.models.ForeignKey', 'models.ForeignKey', (['ModelFileFormat'], {'null': '(True)', 'on_delete': 'models.CASCADE'}), '(ModelFileFormat, null=True, on_delete=models.CASCADE)\n', (2916, 2970), False, 'from django.db import models\n'), ((6146, 6170), 'genui.utils.extensions.tasks.models.PolymorphicTaskManager', 'PolymorphicTaskManager', ([], {}), '()\n', (6168, 6170), False, 'from genui.utils.extensions.tasks.models import TaskShortcutsMixIn, PolymorphicTaskManager\n'), ((6185, 6254), 'django.db.models.ForeignKey', 'models.ForeignKey', (['ModelBuilder'], {'on_delete': 'models.CASCADE', 'null': '(False)'}), '(ModelBuilder, on_delete=models.CASCADE, null=False)\n', (6202, 6254), False, 'from django.db import models\n'), ((7945, 8011), 'django.db.models.ForeignKey', 'models.ForeignKey', (['Algorithm'], {'on_delete': 'models.CASCADE', 'null': '(False)'}), '(Algorithm, on_delete=models.CASCADE, null=False)\n', (7962, 8011), False, 'from django.db import models\n'), ((8023, 8093), 'django.db.models.ForeignKey', 'models.ForeignKey', (['AlgorithmMode'], {'on_delete': 'models.CASCADE', 'null': '(False)'}), '(AlgorithmMode, on_delete=models.CASCADE, null=False)\n', (8040, 8093), False, 'from django.db import models\n'), ((8114, 8216), 'django.db.models.ForeignKey', 'models.ForeignKey', (['Model'], {'null': '(False)', 'on_delete': 'models.CASCADE', 'related_name': '"""trainingStrategies"""'}), "(Model, null=False, on_delete=models.CASCADE, related_name\n    ='trainingStrategies')\n", (8131, 8216), False, 'from django.db import models\n'), ((8275, 8346), 'django.db.models.ForeignKey', 'models.ForeignKey', (['ModelParameter'], {'on_delete': 'models.CASCADE', 'null': '(False)'}), '(ModelParameter, on_delete=models.CASCADE, null=False)\n', (8292, 8346), False, 'from django.db import models\n'), ((8362, 8475), 'django.db.models.ForeignKey', 'models.ForeignKey', (['TrainingStrategy'], {'on_delete': 'NON_POLYMORPHIC_CASCADE', 'null': '(True)', 'related_name': '"""parameters"""'}), "(TrainingStrategy, on_delete=NON_POLYMORPHIC_CASCADE, null\n    =True, related_name='parameters')\n", (8379, 8475), False, 'from django.db import models\n'), ((8599, 8632), 'django.db.models.CharField', 'models.CharField', ([], {'max_length': '(1024)'}), '(max_length=1024)\n', (8615, 8632), False, 'from django.db import models\n'), ((8694, 8725), 'django.db.models.BooleanField', 'models.BooleanField', ([], {'null': '(False)'}), '(null=False)\n', (8713, 8725), False, 'from django.db import models\n'), ((8855, 8886), 'django.db.models.IntegerField', 'models.IntegerField', ([], {'null': '(False)'}), '(null=False)\n', (8874, 8886), False, 'from django.db import models\n'), ((9017, 9046), 'django.db.models.FloatField', 'models.FloatField', ([], {'null': '(False)'}), '(null=False)\n', (9034, 9046), False, 'from django.db import models\n'), ((9412, 9470), 'django.db.models.CharField', 'models.CharField', ([], {'unique': '(True)', 'blank': '(False)', 'max_length': '(128)'}), '(unique=True, blank=False, max_length=128)\n', (9428, 9470), False, 'from django.db import models\n'), ((9488, 9549), 'django.db.models.ManyToManyField', 'models.ManyToManyField', (['AlgorithmMode'], {'related_name': '"""metrics"""'}), "(AlgorithmMode, related_name='metrics')\n", (9510, 9549), False, 'from django.db import models\n'), ((9572, 9629), 'django.db.models.ManyToManyField', 'models.ManyToManyField', (['Algorithm'], {'related_name': '"""metrics"""'}), "(Algorithm, related_name='metrics')\n", (9594, 9629), False, 'from django.db import models\n'), ((9648, 9694), 'django.db.models.TextField', 'models.TextField', ([], {'max_length': '(10000)', 'blank': '(True)'}), '(max_length=10000, blank=True)\n', (9664, 9694), False, 'from django.db import models\n'), ((9850, 9896), 'django.db.models.ManyToManyField', 'models.ManyToManyField', (['ModelPerformanceMetric'], {}), '(ModelPerformanceMetric)\n', (9872, 9896), False, 'from django.db import models\n'), ((9917, 10021), 'django.db.models.ForeignKey', 'models.ForeignKey', (['Model'], {'null': '(False)', 'on_delete': 'models.CASCADE', 'related_name': '"""validationStrategies"""'}), "(Model, null=False, on_delete=models.CASCADE, related_name\n    ='validationStrategies')\n", (9934, 10021), False, 'from django.db import models\n'), ((10063, 10095), 'django.db.models.IntegerField', 'models.IntegerField', ([], {'blank': '(False)'}), '(blank=False)\n', (10082, 10095), False, 'from django.db import models\n'), ((10199, 10229), 'django.db.models.FloatField', 'models.FloatField', ([], {'blank': '(False)'}), '(blank=False)\n', (10216, 10229), False, 'from django.db import models\n'), ((10389, 10468), 'django.db.models.ForeignKey', 'models.ForeignKey', (['ModelPerformanceMetric'], {'null': '(False)', 'on_delete': 'models.CASCADE'}), '(ModelPerformanceMetric, null=False, on_delete=models.CASCADE)\n', (10406, 10468), False, 'from django.db import models\n'), ((10481, 10511), 'django.db.models.FloatField', 'models.FloatField', ([], {'blank': '(False)'}), '(blank=False)\n', (10498, 10511), False, 'from django.db import models\n'), ((10524, 10627), 'django.db.models.ForeignKey', 'models.ForeignKey', (['Model'], {'null': '(False)', 'on_delete': 'NON_POLYMORPHIC_CASCADE', 'related_name': '"""performance"""'}), "(Model, null=False, on_delete=NON_POLYMORPHIC_CASCADE,\n    related_name='performance')\n", (10541, 10627), False, 'from django.db import models\n'), ((10681, 10713), 'django.db.models.IntegerField', 'models.IntegerField', ([], {'blank': '(False)'}), '(blank=False)\n', (10700, 10713), False, 'from django.db import models\n'), ((10770, 10814), 'django.db.models.IntegerField', 'models.IntegerField', ([], {'null': '(False)', 'blank': '(False)'}), '(null=False, blank=False)\n', (10789, 10814), False, 'from django.db import models\n'), ((10826, 10870), 'django.db.models.IntegerField', 'models.IntegerField', ([], {'null': '(False)', 'blank': '(False)'}), '(null=False, blank=False)\n', (10845, 10870), False, 'from django.db import models\n'), ((10922, 10952), 'django.db.models.FloatField', 'models.FloatField', ([], {'blank': '(False)'}), '(blank=False)\n', (10939, 10952), False, 'from django.db import models\n'), ((10963, 11073), 'django.db.models.ForeignKey', 'models.ForeignKey', (['ModelPerformance'], {'null': '(False)', 'on_delete': 'NON_POLYMORPHIC_CASCADE', 'related_name': '"""points"""'}), "(ModelPerformance, null=False, on_delete=\n    NON_POLYMORPHIC_CASCADE, related_name='points')\n", (10980, 11073), False, 'from django.db import models\n'), ((3039, 3057), 'genui.utils.models.OverwriteStorage', 'OverwriteStorage', ([], {}), '()\n', (3055, 3057), False, 'from genui.utils.models import NON_POLYMORPHIC_CASCADE, OverwriteStorage\n'), ((3377, 3389), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (3387, 3389), False, 'import uuid\n'), ((3588, 3600), 'uuid.uuid4', 'uuid.uuid4', ([], {}), '()\n', (3598, 3600), False, 'import uuid\n'), ((4426, 4464), 'os.path.basename', 'os.path.basename', (['model.modelFile.path'], {}), '(model.modelFile.path)\n', (4442, 4464), False, 'import os\n')]