import os
import re
import sys

from qgis.PyQt.QtCore import QCoreApplication
from qgis.core import (
    QgsProcessingAlgorithm,
    QgsProcessingParameterFile, QgsProcessingParameterNumber, QgsProcessingParameterFolderDestination,
    QgsProcessingParameterBoolean, QgsProcessingParameterString, QgsProcessingParameterDefinition,
    QgsProcessingParameterEnum
)


def best_ckpt_path(checkpoint_dir):
    pattern = re.compile(r'val_iou_(\d+\.\d{4})')
    return max(
        (os.path.join(checkpoint_dir, f) for f in os.listdir(checkpoint_dir) if pattern.search(f)),
        key=lambda f: float(pattern.search(f).group(1))
    )


class NullIO:
    def write(self, *args, **kwargs):
        pass

    def flush(self, *args, **kwargs):
        pass


class DL_Trainer(QgsProcessingAlgorithm):
    """DL_Train
    This is an example algorithm that takes a vector layer and
    creates a new identical one.

    It is meant to be used as an example of how to create your own
    algorithms and explain methods and variables used to do it. An
    algorithm like this will be available in all elements, and there
    is not need for additional work.

    All Processing algorithms should extend the QgsProcessingAlgorithm
    class.
    """

    # Constants used to refer to parameters and outputs. They will be
    # used when calling the algorithm from another algorithm, or when
    # calling from the QGIS console.

    train_val_input_folder = 'input_folder'
    arch = 'arch'
    backbone = 'Backbone'
    pretrained_weights = 'pretrained_weights'
    freeze_encoder = 'freeze_encoder'
    data_aug = 'data_aug'
    batch_size = 'Batchsize'
    n_epochs = 'Epochs'
    lr = 'learning rate'
    lr_finder = 'lr_finder'
    pat = 'patience'
    ignore_index = 'ignore index'
    class_weights_balanced = ' class_weights_balanced'
    normalization_flag = 'normalization'
    num_workers = 'Number of workers'
    device = 'Device'
    device_numbers = 'Device Numbers'
    logdirpath = 'logdirpath '
    checkpoint = 'checkpoint'
    n_classes = 'n_classes'
    num_models = 'num_models'
    logdirpath_model = 'logdirpath_model'

    # print_detail_log = 'print_detail_log'

    def tr(self, string):
        """
        Returns a translatable string with the self.tr() function.
        """
        return QCoreApplication.translate('Processing', string)

    def name(self):
        """
        Returns the algorithm name, used for identifying the algorithm. This
        string should be fixed for the algorithm, and must not be localised.
        The name should be unique within each provider. Names should contain
        lowercase alphanumeric characters only and no spaces or other
        formatting characters.
        """
        return 'Deep Learning Trainer'

    def displayName(self):
        """
        Returns the translated algorithm name, which should be used for any
        user-visible display of the algorithm name.
        """
        return self.tr('Deep Learning Trainer')

    def group(self):
        """
        Returns the name of the group this algorithm belongs to. This string
        should be localised.
        """
        return self.tr('SpecDeepMap')

    def groupId(self):
        """
        Returns the unique ID of the group this algorithm belongs to. This
        string should be fixed for the algorithm, and must not be localised.
        The group id should be unique within each provider. Group id should
        contain lowercase alphanumeric characters only and no spaces or other
        formatting characters.
        """
        return 'SpecDeepMap'

    def shortHelpString(self):
        html = (
            ''
            '<p>This algorithm trains a deep learning model specifically a U-Net, U-Net++ or DeeplabV3+, '
            'SegFormer model with a definable backbone for a semantic segmentation task. '
            'For Training the model a Cross-entropy loss or balanced Cross-entropy loss is used. '
            'The performance metric, which is used for training is Intersection over Union.'
            '</p>'
            '<h3>Input folder (Training and validation dataset)</h3>'
            '<p>The input folder should be the one generated by the Dataset Maker, it needs to have a train, '
            'validation and test dataset csv file, as well as a summary csv file and '
            'optionally also a normalization csv file.'
            '</p>'
            '<h3>Deep learning model: architecture</h3>'
            '<p>The main architecture for the deep learning model can be defined with this variable. '
            'Options are U-Net, U-Net++ and DeeplabV3+, SegFormer, these models need to be combine with a '
            'specified backbone. Also 2D-Justo-UNet-Simple can be used as architecture, '
            'but this architecture can not be used with backbone.</p>'
            '<h3>Deep learning model: backbone</h3>'
            '<p> The model can be built with several backbones from pytorch segmentation models library, '
            'e.g. ResNet 50, ResNeXt, EfficientNet, by passing the correct model encoder name as string '
            'into the parameter (so copy name and paste) e.g. by passing the string resnet50 for ResNet-50 or '
            'the string efficientnet-b0  for a small EffitientNet. For available options look into pytorch '
            'segmentation models libary here https://smp.readthedocs.io/en/latest/encoders.html . '
            'Also a range of Timm Encoder can also be used, like advanced models ConvNext or Swin-Transformer. '
            'For using these models also past the encoder string into the parameter of the model with the prefix  '
            'tu- followed by timm-encoder name  e.g. tu-convnext_tiny  for a small ConvNext. '
            'You can find list of available Timm encoders here: '
            'https://smp.readthedocs.io/en/latest/encoders_timm.html '
            '</p>'
            '<h3>Load pretrained weights </h3>'
            '<p>Here a user can chose weights for model initialization. The options are None for training from '
            'scratch and pretrained imagenet weights for RGB images which can be extended to more spectral bands. '
            'The input data need to be structured in bgr + additional channels format for imagenet weights use as '
            'well as requires data normalization. Further pretrained weights for  Sentinel-2 Top of Atmosphere '
            'are provided for ResNet-18 and ResNet-50. These weights stem from foundation model training from '
            'SSL4EO-S12 by Wang et al. 2023. The automatic preprocessing for these weights is scaling the input '
            'data by 10000, any other scaling or normalization will be ignored if these weights are chosen.'
            '</p>'
            '<h3>Load model from path</h3>'
            '<p>Load model from path for continuing training or to initialize a model. Restriction is that model must '
            'be compatible with training scheme and input data type.'
            '</p>'
            '<h3>Freeze backbone </h3>'
            '<p>This freezes the weight of the backbone for training. This is a routine step for '
            'transferlearning where in the first step a model is trained on new classes with a frozen backbone '
            'and in a second training step it is fine-tuned with an unfrozen backbone. To achieve this here, '
            'train first with frozen backbone and then load trained model again and train with unfrozen backbone.'
            '</p>'
            '<h3>Data augmentation (random flip ) </h3>'
            '<p>This apply data augmentation. Random vertical and horizontal flip. Each augmentation has a '
            'probability of occurring of 50 %. This data augmentation is happening on the fly and '
            'prevents overfitting of the model.'
            '</p>'
            '<h3>Early stopping</h3>'
            '<p>This stops the model when validation loss is not improving for 50 epochs.</p>'
            '<h3>Balanced Training using Class Weights</h3>'
            '<p>This parameter enables balanced training, for this the precomputed class weights based on the '
            'training dataset are used, which are listed in the summary csv file.'
            '</p>'
            '<h3>Data Normalization</h3>'
            '<p>This parameter normalizes the image data with mean and std. per channel. To use this parameter '
            'the normalization statistic had to be created using the Dataset Maker.'
            '</p>'
            '<h3>Batch size</h3>'
            '<p>This defines the number of images which are processed in batches. </p>'
            '<h3>Epochs</h3>'
            '<p>This defines the number of Epochs which are used to train a model. One epochs means the model is '
            'trained once on the whole training dataset. Epochs are 0 indexed so the first training epoch is epoch-0. '
            '</p>'
            '<h3>Learning rate</h3>'
            '<p>This defines the Learning rate for the Adam optimizer of the model. Learning rate is decreasing '
            'during training using CosineAnnealingLR using defined epochs value as T_max value, so lr decreases '
            'steadily towards near zero during whole training run (min lr: 1e-6). '
            '(More detail see here: '
            'https://docs.pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html )'
            '</p>'
            '<h3>Automatic learning rate finder</h3>'
            '<p>If activated this parameter runs a learning rate finder, meaning it test 100 learning rates '
            'before training and picks the one with the most stable learning for training. If chosen this '
            'overwrites the manual defined learning rate. The found and implemented learning-rate is printed '
            'to log interface. '
            '</p>'
            '<h3>Number of workers</h3>'
            '<p>This defines number of cpus used for data loading and augmentation and supports your training speed.'
            '</p>'
            '<h3>Type of device</h3>'
            '<p> You can use either CPU or GPU for training. If available use GPU.</p>'
            '<h3>Number of devices</h3>'
            '<p> For distributated training you can also here define how many GPUs you want to use.</p>'
            '<h3>Number of models</h3>'
            '<p> Defines how many models should be saved. -1 means each epoch a model is saved. '
            'Other integer define a limited number of models.</p>'
            '<h3>Path for saving Tensorboard logger</h3>'
            '<p>Define folder where Tensorboard logger is saved.</p>'
            '<h3>Path for saving model</h3>'
            '<p>Define folder where models are saved</p>'
        )

        return html

    def initAlgorithm(self, config=None):
        """
        Here we define the inputs and output of the algorithm, along
        with some other properties.
        """

        self.addParameter(QgsProcessingParameterFile(
            name=self.train_val_input_folder, description='Input folder (Train and Validation dataset)',
            behavior=QgsProcessingParameterFile.Behavior.Folder))
        self.addParameter(QgsProcessingParameterEnum(
            name=self.arch, description='Model architecture',
            options=['U-Net', 'U-Net++', 'DeepLabV3+', 'SegFormer', '2D-Justo-UNet-Simple'], defaultValue=0))
        self.addParameter(QgsProcessingParameterString(
            name=self.backbone, description='Model backbone', defaultValue='resnet18'))
        self.addParameter(QgsProcessingParameterEnum(
            name=self.pretrained_weights, description='Load pretrained weights',
            options=['imagenet', 'None', 'Sentinel_2_TOA_Resnet18', 'Sentinel_2_TOA_Resnet50'], defaultValue=1))
        self.addParameter(
            QgsProcessingParameterFile(self.checkpoint, description='Load model from path', optional=True))
        self.addParameter(QgsProcessingParameterBoolean(
            name=self.freeze_encoder, description='Freeze backbone',
            defaultValue=True))
        self.addParameter(QgsProcessingParameterBoolean(
            name=self.data_aug, description='Data augmentation (random flip & rotate by 45°)',
            defaultValue=True))
        self.addParameter(QgsProcessingParameterBoolean(
            name=self.pat, description='Early stopping',
            defaultValue=True))
        self.addParameter(
            QgsProcessingParameterBoolean(self.class_weights_balanced, self.tr('Balanced Training using Class Weights'),
                                          defaultValue=True))
        self.addParameter(
            QgsProcessingParameterBoolean(self.normalization_flag, self.tr('Data Normalization'), defaultValue=False))
        self.addParameter(QgsProcessingParameterNumber(
            name=self.batch_size, description='Batch size', type=QgsProcessingParameterNumber.Integer,
            defaultValue=2, minValue=1))
        self.addParameter(QgsProcessingParameterNumber(
            name=self.n_epochs, description='Epochs', type=QgsProcessingParameterNumber.Integer,
            defaultValue=1, minValue=1))

        self.addParameter(QgsProcessingParameterNumber(
            name=self.lr, description='Learning rate', type=QgsProcessingParameterNumber.Double,
            defaultValue=0.001, minValue=0.0000001))
        self.addParameter(QgsProcessingParameterBoolean(
            name=self.lr_finder, description='Automatic learning rate finder',
            defaultValue=False))

        self.addParameter(QgsProcessingParameterEnum(
            name=self.device, description='Type of device (GPU/CPU)', options=['cpu', 'gpu'], defaultValue=0))

        # adjusted   to make advanced

        p = QgsProcessingParameterNumber(
            name=self.num_workers, description='Number of workers (CPUs used for data loading and augumenation)',
            type=QgsProcessingParameterNumber.Integer,
            defaultValue=0)
        p.setFlags(p.flags() | QgsProcessingParameterDefinition.Flag.FlagAdvanced)
        self.addParameter(p)

        p1 = QgsProcessingParameterNumber(
            name=self.device_numbers, description='Number of devices (GPU/CPU)', optional=True,
            type=QgsProcessingParameterNumber.Integer,
            defaultValue=1, minValue=0)
        p1.setFlags(p1.flags() | QgsProcessingParameterDefinition.Flag.FlagAdvanced)
        self.addParameter(p1)

        p2 = QgsProcessingParameterNumber(
            name=self.num_models, description='Number of Models', optional=True,
            type=QgsProcessingParameterNumber.Integer,
            defaultValue=-1)
        p2.setFlags(p2.flags() | QgsProcessingParameterDefinition.Flag.FlagAdvanced)
        self.addParameter(p2)

        # p3 = QgsProcessingParameterBoolean(
        #   name=self.print_detail_log, description='Print detail train process in python console', optional=True,
        #    defaultValue=False)
        # p3.setFlags(p3.flags() | QgsProcessingParameterDefinition.Flag.FlagAdvanced)
        # self.addParameter(p3)

        self.addParameter(QgsProcessingParameterFolderDestination(
            name=self.logdirpath, description='Path for saving Tensorboard logger'))

        self.addParameter(QgsProcessingParameterFolderDestination(
            name=self.logdirpath_model, description='Path for saving model'))

    def processAlgorithm(self, parameters, context, feedback):
        """
        Here is where the processing itself takes place.
        """
        # control print out in python consol or non
        # print_detail =self.parameterAsBool(parameters, self.print_detail_log, context)

        # if print_detail==False:
        # Save original stdout and stderr
        original_stdout = sys.stdout
        original_stderr = sys.stderr

        # Redirect stdout and stderr to null handler to suppress output
        sys.stdout = NullIO()
        sys.stderr = NullIO()

        # main function
        from enmapbox.apps.SpecDeepMap.core_deep_learning_trainer import dl_train

        feedback.pushInfo(
            "If you use pretrained weights of specific backbone for the first time, "
            "they will be downloaded before the training starts.")

        _ = dl_train(
            input_folder=self.parameterAsString(parameters, self.train_val_input_folder, context),
            arch_index=self.parameterAsEnum(parameters, self.arch, context),
            backbone=self.parameterAsString(parameters, self.backbone, context),
            pretrained_weights_index=self.parameterAsEnum(parameters, self.pretrained_weights, context),
            checkpoint_path=self.parameterAsFile(parameters, self.checkpoint, context),
            freeze_encoder=self.parameterAsBool(parameters, self.freeze_encoder, context),
            data_aug=self.parameterAsBool(parameters, self.data_aug, context),
            batch_size=self.parameterAsInt(parameters, self.batch_size, context),
            n_epochs=self.parameterAsInt(parameters, self.n_epochs, context),
            lr=self.parameterAsDouble(parameters, self.lr, context),
            tune=self.parameterAsBool(parameters, self.lr_finder, context),
            early_stop=self.parameterAsBool(parameters, self.pat, context),
            class_weights_balanced=self.parameterAsBool(parameters, self.class_weights_balanced, context),
            normalization_bool=self.parameterAsBool(parameters, self.normalization_flag, context),
            num_workers=self.parameterAsInt(parameters, self.num_workers, context),
            num_models=self.parameterAsInt(parameters, self.num_models, context),
            acc_type_index=self.parameterAsEnum(parameters, self.device, context),
            acc_type_numbers=self.parameterAsInt(parameters, self.device_numbers, context),
            logdirpath=self.parameterAsString(parameters, self.logdirpath, context),
            logdirpath_model=self.parameterAsString(parameters, self.logdirpath_model, context),
            feedback=feedback)

        feedback.pushInfo("Training completed.")

        out = self.parameterAsString(parameters, self.logdirpath, context)
        out_m = self.parameterAsString(parameters, self.logdirpath_model, context)
        #
        # select best iou model automatic so can be used in model builder
        best_iou_model = best_ckpt_path(out_m)

        outputs = {self.logdirpath: out, self.logdirpath_model: best_iou_model}

        # set variable print variable back
        # if print_detail==False:
        # Restore original stdout and stderr
        sys.stdout = original_stdout
        sys.stderr = original_stderr
        # Ensure flush of any buffered output
        # sys.stdout.flush()
        # sys.stderr.flush()

        return outputs

    # 6
    def helpUrl(self, *args, **kwargs):
        return ''

    # 7
    def createInstance(self):
        return type(self)()
