#!/usr/bin/env python

# Copyright 2015-2020 Earth Sciences Department, BSC-CNS

# This file is part of Autosubmit.

# Autosubmit 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.

# Autosubmit 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 Autosubmit.  If not, see <http://www.gnu.org/licenses/>.

from autosubmit.platforms.wrappers.wrapper_builder import WrapperDirector, PythonVerticalWrapperBuilder, \
    PythonHorizontalWrapperBuilder, PythonHorizontalVerticalWrapperBuilder, PythonVerticalHorizontalWrapperBuilder, \
    BashHorizontalWrapperBuilder, BashVerticalWrapperBuilder, SrunHorizontalWrapperBuilder,SrunVerticalHorizontalWrapperBuilder
from autosubmit.config.config_common import AutosubmitConfig


class WrapperFactory(object):

    def __init__(self, platform):
        self.platform = platform
        self.wrapper_director = WrapperDirector()
        self.exception = "This type of wrapper is not supported for this platform"
        self.configuration = None
    def get_wrapper(self, wrapper_builder, **kwargs):
        kwargs['allocated_nodes'] = self.allocated_nodes()
        kwargs['dependency'] = self.dependency(kwargs['dependency'])
        kwargs['queue'] = self.queue(kwargs['queue'])
        kwargs['header_directive'] = self.header_directives(**kwargs)
        kwargs['configuration'] = self.configuration
        kwargs['language'] = self.configuration.get_wrapper_language()
        builder = wrapper_builder(**kwargs)
        return self.wrapper_director.construct(builder)

    def vertical_wrapper(self, **kwargs):
        raise NotImplemented(self.exception)

    def horizontal_wrapper(self, **kwargs):
        raise NotImplemented(self.exception)

    def hybrid_wrapper_horizontal_vertical(self, **kwargs):
        raise NotImplemented(self.exception)

    def hybrid_wrapper_vertical_horizontal(self, **kwargs):
        raise NotImplemented(self.exception)

    def header_directives(self, **kwargs):
        pass

    def allocated_nodes(self):
        return ''

    def dependency(self, dependency):
        return '#' if dependency is None else self.dependency_directive(dependency)

    def queue(self, queue):
        return '#' if not queue else self.queue_directive(queue)

    def dependency_directive(self, dependency):
        pass

    def queue_directive(self, queue):
        pass


class SlurmWrapperFactory(WrapperFactory):

    def vertical_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()
        return PythonVerticalWrapperBuilder(**kwargs)

    def horizontal_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        if kwargs["method"] == 'srun':
            return SrunHorizontalWrapperBuilder(**kwargs)
        else:
            return PythonHorizontalWrapperBuilder(**kwargs)

    def hybrid_wrapper_horizontal_vertical(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return PythonHorizontalVerticalWrapperBuilder(**kwargs)

    def hybrid_wrapper_vertical_horizontal(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        if kwargs["method"] == 'srun':
            return SrunVerticalHorizontalWrapperBuilder(**kwargs)
        else:
            return PythonVerticalHorizontalWrapperBuilder(**kwargs)

    def header_directives(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()
        return self.platform.wrapper_header(kwargs['name'], kwargs['queue'], kwargs['project'], kwargs['wallclock'],
                                            kwargs['num_processors'], kwargs['dependency'], kwargs['directives'],kwargs['threads'],kwargs['method'],kwargs['language'])

    def allocated_nodes(self):
        return self.platform.allocated_nodes()

    def dependency_directive(self, dependency):
        return '#SBATCH --dependency=afterok:{0}'.format(dependency)

    def queue_directive(self, queue):
        return '#SBATCH --qos={0}'.format(queue)

class PJMWrapperFactory(WrapperFactory):

    def vertical_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()
        return PythonVerticalWrapperBuilder(**kwargs)

    def horizontal_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        if kwargs["method"] == 'srun':
            return SrunHorizontalWrapperBuilder(**kwargs)
        else:
            return PythonHorizontalWrapperBuilder(**kwargs)

    def hybrid_wrapper_horizontal_vertical(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return PythonHorizontalVerticalWrapperBuilder(**kwargs)

    def hybrid_wrapper_vertical_horizontal(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        if kwargs["method"] == 'srun':
            return SrunVerticalHorizontalWrapperBuilder(**kwargs)
        else:
            return PythonVerticalHorizontalWrapperBuilder(**kwargs)

    def header_directives(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return self.platform.wrapper_header(kwargs)

    def allocated_nodes(self):
        return self.platform.allocated_nodes()


    def queue_directive(self, queue):
        return '#PJM --qos={0}'.format(queue)
    def partition_directive(self, partition):
        return '#PJM --partition={0}'.format(partition)
    def exclusive_directive(self, exclusive):
        return '#PJM --exclusive'
    def tasks_directive(self, tasks):
        return "#PJM --mpi max-proc-per-node={0}".format(tasks) # searchhint
    def nodes_directive(self, nodes):
        return '#PJM -N {0}'.format(nodes)
    def processors_directive(self, processors):
        return '#PJM -n {0}'.format(processors)
    def threads_directive(self, threads):
        return "export OMP_NUM_THREADS={threads}".format(threads=threads)

    def queue_directive(self, queue):
        return '#PJM -L rscgrp={0}'.format(queue)

    def partition_directive(self, partition):
        return '#PJM -g {0}'.format(partition)

class LSFWrapperFactory(WrapperFactory):

    def vertical_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return PythonVerticalWrapperBuilder(**kwargs)

    def horizontal_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return PythonHorizontalWrapperBuilder(**kwargs)

    def header_directives(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return self.platform.wrapper_header(kwargs['name'], kwargs['queue'], kwargs['project'], kwargs['wallclock'],
                                            kwargs['num_processors'], kwargs['dependency'], kwargs['directives'],kwargs['threads'],kwargs['language'])

    def queue_directive(self, queue):
        return queue

    def dependency_directive(self, dependency):
        return '#BSUB -w \'done({0})\' [-ti]'.format(dependency)


class EcWrapperFactory(WrapperFactory):

    def vertical_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return BashVerticalWrapperBuilder(**kwargs)

    def horizontal_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return BashHorizontalWrapperBuilder(**kwargs)

    def header_directives(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return self.platform.wrapper_header(kwargs['name'], kwargs['queue'], kwargs['project'], kwargs['wallclock'],
                                            kwargs['num_processors'], kwargs['expid'], kwargs['dependency'],
                                            kwargs['rootdir'], kwargs['directives'],kwargs['threads'])

    def queue_directive(self, queue):
        return queue

    def dependency_directive(self, dependency):
        return '#PBS -v depend=afterok:{0}'.format(dependency)


    def horizontal_wrapper(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        if kwargs["method"] == 'srun':
            return SrunHorizontalWrapperBuilder(**kwargs)
        else:
            return PythonHorizontalWrapperBuilder(**kwargs)

    def hybrid_wrapper_horizontal_vertical(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        return PythonHorizontalVerticalWrapperBuilder(**kwargs)

    def hybrid_wrapper_vertical_horizontal(self, **kwargs):
        kwargs['language'] = self.configuration.get_wrapper_language()

        if kwargs["method"] == 'srun':
            return SrunVerticalHorizontalWrapperBuilder(**kwargs)
        else:
            return PythonVerticalHorizontalWrapperBuilder(**kwargs)

    def header_directives(self, **kwargs):
        return self.platform.wrapper_header(kwargs['name'], kwargs['queue'], kwargs['project'], kwargs['wallclock'],
                                            kwargs['num_processors'], kwargs['dependency'], kwargs['directives'],kwargs['threads'],kwargs['method'],kwargs['partition'],self.configuration.get_wrapper_language())

    def allocated_nodes(self):
        return self.platform.allocated_nodes()

    #def dependency_directive(self, dependency):
    #    # There is no option for afterok in the PJM scheduler, but I think it is not needed.
    #    return '#PJM --dependency=afterok:{0}'.format(dependency)

    def queue_directive(self, queue):
        return '#PJM -L rscgrp={0}'.format(queue)

    def partition_directive(self, partition):
        return '#PJM -g {0}'.format(partition)
