# -*- coding: utf-8 -*-
"""
    pygments.lexers.haskell
    ~~~~~~~~~~~~~~~~~~~~~~~

    Lexers for Haskell and related languages.

    :copyright: Copyright 2006-2017 by the Pygments team, see AUTHORS.
    :license: BSD, see LICENSE for details.
"""

import re

from pygments.lexer import Lexer, RegexLexer, bygroups, do_insertions, \
    default, include
from pygments.token import Text, Comment, Operator, Keyword, Name, String, \
    Number, Punctuation, Generic
from pygments import unistring as uni

__all__ = ['HaskellLexer', 'IdrisLexer', 'AgdaLexer', 'CryptolLexer',
           'LiterateHaskellLexer', 'LiterateIdrisLexer', 'LiterateAgdaLexer',
           'LiterateCryptolLexer', 'KokaLexer']

line_re = re.compile('.*?\n')

class HaskellLexer(RegexLexer):
    """
    A Haskell lexer based on the lexemes defined in the Haskell 98 Report.

    .. versionadded:: 0.8
    """
    name = 'Haskell'
    aliases = ['haskell', 'hs']
    filenames = ['*.hs']
    mimetypes = ['text/x-haskell']

    flags = re.MULTILINE | re.UNICODE

    reserved = ('case', 'class', 'data', 'default', 'deriving', 'do', 'else',
                'family', 'if', 'in', 'infix[lr]?', 'instance',
                'let', 'newtype', 'of', 'then', 'type', 'where', '_')
    ascii = ('NUL', 'SOH', '[SE]TX', 'EOT', 'ENQ', 'ACK',
             'BEL', 'BS', 'HT', 'LF', 'VT', 'FF', 'CR', 'S[OI]', 'DLE',
             'DC[1-4]', 'NAK', 'SYN', 'ETB', 'CAN',
             'EM', 'SUB', 'ESC', '[FGRU]S', 'SP', 'DEL')

    tokens = {
        'root': [
            # Whitespace:
            (r'\s+', Text),
            # (r'--\s*|.*$', Comment.Doc),
            (r'--(?![!#$%&*+./<=>?@^|_~:\\]).*?$', Comment.Single),
            (r'\{-', Comment.Multiline, 'comment'),
            # Lexemes:
            #  Identifiers
            (r'\bimport\b', Keyword.Reserved, 'import'),
            (r'\bmodule\b', Keyword.Reserved, 'module'),
            (r'\berror\b', Name.Exception),
            (r'\b(%s)(?!\')\b' % '|'.join(reserved), Keyword.Reserved),
            (r"'[^\\]'", String.Char),  # this has to come before the TH quote
            (r'^[_' + uni.Ll + r'][\w\']*', Name.Function),
            (r"'?[_" + uni.Ll + r"][\w']*", Name),
            (r"('')?[" + uni.Lu + r"][\w\']*", Keyword.Type),
            (r"(')[" + uni.Lu + r"][\w\']*", Keyword.Type),
            (r"(')\[[^\]]*\]", Keyword.Type),  # tuples and lists get special treatment in GHC
            (r"(')\([^)]*\)", Keyword.Type),  # ..
            #  Operators
            (r'\\(?![:!#$%&*+.\\/<=>?@^|~-]+)', Name.Function),  # lambda operator
            (r'(<-|::|->|=>|=)(?![:!#$%&*+.\\/<=>?@^|~-]+)', Operator.Word),  # specials
            (r':[:!#$%&*+.\\/<=>?@^|~-]*', Keyword.Type),  # Constructor operators
            (r'[:!#$%&*+.\\/<=>?@^|~-]+', Operator),  # Other operators
            #  Numbers
            (r'\d+[eE][+-]?\d+', Number.Float),
            (r'\d+\.\d+([eE][+-]?\d+)?', Number.Float),
            (r'0[oO][0-7]+', Number.Oct),
            (r'0[xX][\da-fA-F]+', Number.Hex),
            (r'\d+', Number.Integer),
            #  Character/String Literals
            (r"'", String.Char, 'character'),
            (r'"', String, 'string'),
            #  Special
            (r'\[\]', Keyword.Type),
            (r'\(\)', Name.Builtin),
            (r'[][(),;`{}]', Punctuation),
        ],
        'import': [
            # Import statements
            (r'\s+', Text),
            (r'"', String, 'string'),
            # after "funclist" state
            (r'\)', Punctuation, '#pop'),
            (r'qualified\b', Keyword),
            # import X as Y
            (r'([' + uni.Lu + r'][\w.]*)(\s+)(as)(\s+)([' + uni.Lu + r'][\w.]*)',
             bygroups(Name.Namespace, Text, Keyword, Text, Name), '#pop'),
            # import X hiding (functions)
            (r'([' + uni.Lu + r'][\w.]*)(\s+)(hiding)(\s+)(\()',
             bygroups(Name.Namespace, Text, Keyword, Text, Punctuation), 'funclist'),
            # import X (functions)
            (r'([' + uni.Lu + r'][\w.]*)(\s+)(\()',
             bygroups(Name.Namespace, Text, Punctuation), 'funclist'),
            # import X
            (r'[\w.]+', Name.Namespace, '#pop'),
        ],
        'module': [
            (r'\s+', Text),
            (r'([' + uni.Lu + r'][\w.]*)(\s+)(\()',
             bygroups(Name.Namespace, Text, Punctuation), 'funclist'),
            (r'[' + uni.Lu + r'][\w.]*', Name.Namespace, '#pop'),
        ],
        'funclist': [
            (r'\s+', Text),
            (r'[' + uni.Lu + r']\w*', Keyword.Type),
            (r'(_[\w\']+|[' + uni.Ll + r'][\w\']*)', Name.Function),
            (r'--(?![!#$%&*+./<=>?@^|_~:\\]).*?$', Comment.Single),
            (r'\{-', Comment.Multiline, 'comment'),
            (r',', Punctuation),
            (r'[:!#$%&*+.\\/<=>?@^|~-]+', Operator),
            # (HACK, but it makes sense to push two instances, believe me)
            (r'\(', Punctuation, ('funclist', 'funclist')),
            (r'\)', Punctuation, '#pop:2'),
        ],
        # NOTE: the next four states are shared in the AgdaLexer; make sure
        # any change is compatible with Agda as well or copy over and change
        'comment': [
            # Multiline Comments
            (r'[^-{}]+', Comment.Multiline),
            (r'\{-', Comment.Multiline, '#push'),
            (r'-\}', Comment.Multiline, '#pop'),
            (r'[-{}]', Comment.Multiline),
        ],
        'character': [
            # Allows multi-chars, incorrectly.
            (r"[^\\']'", String.Char, '#pop'),
            (r"\\", String.Escape, 'escape'),
            ("'", String.Char, '#pop'),
        ],
        'string': [
            (r'[^\\"]+', String),
            (r"\\", String.Escape, 'escape'),
            ('"', String, '#pop'),
        ],
        'escape': [
            (r'[abfnrtv"\'&\\]', String.Escape, '#pop'),
            (r'\^[][' + uni.Lu + r'@^_]', String.Escape, '#pop'),
            ('|'.join(ascii), String.Escape, '#pop'),
            (r'o[0-7]+', String.Escape, '#pop'),
            (r'x[\da-fA-F]+', String.Escape, '#pop'),
            (r'\d+', String.Escape, '#pop'),
            (r'\s+\\', String.Escape, '#pop'),
        ],
    }

class IdrisLexer(RegexLexer):
    """
    A lexer for the dependently typed programming language Idris.

    Based on the Haskell and Agda Lexer.

    .. versionadded:: 2.0
    """
    name = 'Idris'
    aliases = ['idris', 'idr']
    filenames = ['*.idr']
    mimetypes = ['text/x-idris']

    reserved = ('case', 'class', 'data', 'default', 'using', 'do', 'else',
                'if', 'in', 'infix[lr]?', 'instance', 'rewrite', 'auto',
                'namespace', 'codata', 'mutual', 'private', 'public', 'abstract',
                'total', 'partial',
                'let', 'proof', 'of', 'then', 'static', 'where', '_', 'with',
                'pattern',  'term',  'syntax', 'prefix',
                'postulate', 'parameters', 'record', 'dsl', 'impossible', 'implicit',
                'tactics', 'intros', 'intro', 'compute', 'refine', 'exact', 'trivial')

    ascii = ('NUL', 'SOH', '[SE]TX', 'EOT', 'ENQ', 'ACK',
             'BEL', 'BS', 'HT', 'LF', 'VT', 'FF', 'CR', 'S[OI]', 'DLE',
             'DC[1-4]', 'NAK', 'SYN', 'ETB', 'CAN',
             'EM', 'SUB', 'ESC', '[FGRU]S', 'SP', 'DEL')

    directives = ('lib', 'link', 'flag', 'include', 'hide', 'freeze', 'access',
                  'default', 'logging', 'dynamic', 'name', 'error_handlers', 'language')

    tokens = {
        'root': [
            # Comments
            (r'^(\s*)(%%%s)' % '|'.join(directives),
             bygroups(Text, Keyword.Reserved)),
            (r'(\s*)(--(?![!#$%&*+./<=>?@^|_~:\\]).*?)$', bygroups(Text, Comment.Single)),
            (r'(\s*)(\|{3}.*?)$', bygroups(Text, Comment.Single)),
            (r'(\s*)(\{-)', bygroups(Text, Comment.Multiline), 'comment'),
            # Declaration
            (r'^(\s*)([^\s(){}]+)(\s*)(:)(\s*)',
             bygroups(Text, Name.Function, Text, Operator.Word, Text)),
            #  Identifiers
            (r'\b(%s)(?!\')\b' % '|'.join(reserved), Keyword.Reserved),
            (r'(import|module)(\s+)', bygroups(Keyword.Reserved, Text), 'module'),
            (r"('')?[A-Z][\w\']*", Keyword.Type),
            (r'[a-z][\w\']*', Text),
            #  Special Symbols
            (r'(<-|::|->|=>|=)', Operator.Word),  # specials
            (r'([(){}\[\]:!#$%&*+.\\/<=>?@^|~-]+)', Operator.Word),  # specials
            #  Numbers
            (r'\d+[eE][+-]?\d+', Number.Float),
            (r'\d+\.\d+([eE][+-]?\d+)?', Number.Float),
            (r'0[xX][\da-fA-F]+', Number.Hex),
            (r'\d+', Number.Integer),
            # Strings
            (r"'", String.Char, 'character'),
            (r'"', String, 'string'),
            (r'[^\s(){}]+', Text),
            (r'\s+?', Text),  # Whitespace
        ],
        'module': [
            (r'\s+', Text),
            (r'([A-Z][\w.]*)(\s+)(\()',
             bygroups(Name.Namespace, Text, Punctuation), 'funclist'),
            (r'[A-Z][\w.]*', Name.Namespace, '#pop'),
        ],
        'funclist': [
            (r'\s+', Text),
            (r'[A-Z]\w*', Keyword.Type),
            (r'(_[\w\']+|[a-z][\w\']*)', Name.Function),
            (r'--.*$', Comment.Single),
            (r'\{-', Comment.Multiline, 'comment'),
            (r',', Punctuation),
            (r'[:!#$%&*+.\\/<=>?@^|~-]+', Operator),
            # (HACK, but it makes sense to push two instances, believe me)
            (r'\(', Punctuation, ('funclist', 'funclist')),
            (r'\)', Punctuation, '#pop:2'),
        ],
        # NOTE: the next four states are shared in the AgdaLexer; make sure
        # any change is compatible with Agda as well or copy over and change
        'comment': [
            # Multiline Comments
            (r'[^-{}]+', Comment.Multiline),
            (r'\{-', Comment.Multiline, '#push'),
            (r'-\}', Comment.Multiline, '#pop'),
            (r'[-{}]', Comment.Multiline),
        ],
        'character': [
            # Allows multi-chars, incorrectly.
            (r"[^\\']", String.Char),
            (r"\\", String.Escape, 'escape'),
            ("'", String.Char, '#pop'),
        ],
        'string': [
            (r'[^\\"]+', String),
            (r"\\", String.Escape, 'escape'),
            ('"', String, '#pop'),
        ],
        'escape': [
            (r'[abfnrtv"\'&\\]', String.Escape, '#pop'),
            (r'\^[][A-Z@^_]', String.Escape, '#pop'),
            ('|'.join(ascii), String.Escape, '#pop'),
            (r'o[0-7]+', String.Escape, '#pop'),
            (r'x[\da-fA-F]+', String.Escape, '#pop'),
            (r'\d+', String.Escape, '#pop'),
            (r'\s+\\', String.Escape, '#pop')
        ],
    }

class AgdaLexer(RegexLexer):
    """
    For the `Agda <http://wiki.portal.chalmers.se/agda/pmwiki.php>`_
    dependently typed functional programming language and proof assistant.

    .. versionadded:: 2.0
    """

    name = 'Agda'
    aliases = ['agda']
    filenames = ['*.agda']
    mimetypes = ['text/x-agda']

    reserved = ['abstract', 'codata', 'coinductive', 'constructor', 'data',
                'field', 'forall', 'hiding', 'in', 'inductive', 'infix',
                'infixl', 'infixr', 'instance', 'let', 'mutual', 'open',
                'pattern', 'postulate', 'primitive', 'private',
                'quote', 'quoteGoal', 'quoteTerm',
                'record', 'renaming', 'rewrite', 'syntax', 'tactic',
                'unquote', 'unquoteDecl', 'using', 'where', 'with']

    tokens = {
        'root': [
            # Declaration
            (r'^(\s*)([^\s(){}]+)(\s*)(:)(\s*)',
             bygroups(Text, Name.Function, Text, Operator.Word, Text)),
            # Comments
            (r'--(?![!#$%&*+./<=>?@^|_~:\\]).*?$', Comment.Single),
            (r'\{-', Comment.Multiline, 'comment'),
            # Holes
            (r'\{!', Comment.Directive, 'hole'),
            # Lexemes:
            #  Identifiers
            (r'\b(%s)(?!\')\b' % '|'.join(reserved), Keyword.Reserved),
            (r'(import|module)(\s+)', bygroups(Keyword.Reserved, Text), 'module'),
            (r'\b(Set|Prop)\b', Keyword.Type),
            #  Special Symbols
            (r'(\(|\)|\{|\})', Operator),
            (u'(\\.{1,3}|\\||\u039B|\u2200|\u2192|:|=|->)', Operator.Word),
            #  Numbers
            (r'\d+[eE][+-]?\d+', Number.Float),
            (r'\d+\.\d+([eE][+-]?\d+)?', Number.Float),
            (r'0[xX][\da-fA-F]+', Number.Hex),
            (r'\d+', Number.Integer),
            # Strings
            (r"'", String.Char, 'character'),
            (r'"', String, 'string'),
            (r'[^\s(){}]+', Text),
            (r'\s+?', Text),  # Whitespace
        ],
        'hole': [
            # Holes
            (r'[^!{}]+', Comment.Directive),
            (r'\{!', Comment.Directive, '#push'),
            (r'!\}', Comment.Directive, '#pop'),
            (r'[!{}]', Comment.Directive),
        ],
        'module': [
            (r'\{-', Comment.Multiline, 'comment'),
            (r'[a-zA-Z][\w.]*', Name, '#pop'),
            (r'[\W0-9_]+', Text)
        ],
        'comment': HaskellLexer.tokens['comment'],
        'character': HaskellLexer.tokens['character'],
        'string': HaskellLexer.tokens['string'],
        'escape': HaskellLexer.tokens['escape']
    }

class CryptolLexer(RegexLexer):
    """
    FIXME: A Cryptol2 lexer based on the lexemes defined in the Haskell 98 Report.

    .. versionadded:: 2.0
    """
    name = 'Cryptol'
    aliases = ['cryptol', 'cry']
    filenames = ['*.cry']
    mimetypes = ['text/x-cryptol']

    reserved = ('Arith', 'Bit', 'Cmp', 'False', 'Inf', 'True', 'else',
                'export', 'extern', 'fin', 'if', 'import', 'inf', 'lg2',
                'max', 'min', 'module', 'newtype', 'pragma', 'property',
                'then', 'type', 'where', 'width')
    ascii = ('NUL', 'SOH', '[SE]TX', 'EOT', 'ENQ', 'ACK',
             'BEL', 'BS', 'HT', 'LF', 'VT', 'FF', 'CR', 'S[OI]', 'DLE',
             'DC[1-4]', 'NAK', 'SYN', 'ETB', 'CAN',
             'EM', 'SUB', 'ESC', '[FGRU]S', 'SP', 'DEL')

    tokens = {
        'root': [
            # Whitespace:
            (r'\s+', Text),
            # (r'--\s*|.*$', Comment.Doc),
            (r'//.*$', Comment.Single),
            (r'/\*', Comment.Multiline, 'comment'),
            # Lexemes:
            #  Identifiers
            (r'\bimport\b', Keyword.Reserved, 'import'),
            (r'\bmodule\b', Keyword.Reserved, 'module'),
            (r'\berror\b', Name.Exception),
            (r'\b(%s)(?!\')\b' % '|'.join(reserved), Keyword.Reserved),
            (r'^[_a-z][\w\']*', Name.Function),
            (r"'?[_a-z][\w']*", Name),
            (r"('')?[A-Z][\w\']*", Keyword.Type),
            #  Operators
            (r'\\(?![:!#$%&*+.\\/<=>?@^|~-]+)', Name.Function),  # lambda operator
            (r'(<-|::|->|=>|=)(?![:!#$%&*+.\\/<=>?@^|~-]+)', Operator.Word),  # specials
            (r':[:!#$%&*+.\\/<=>?@^|~-]*', Keyword.Type),  # Constructor operators
            (r'[:!#$%&*+.\\/<=>?@^|~-]+', Operator),  # Other operators
            #  Numbers
            (r'\d+[eE][+-]?\d+', Number.Float),
            (r'\d+\.\d+([eE][+-]?\d+)?', Number.Float),
            (r'0[oO][0-7]+', Number.Oct),
            (r'0[xX][\da-fA-F]+', Number.Hex),
            (r'\d+', Number.Integer),
            #  Character/String Literals
            (r"'", String.Char, 'character'),
            (r'"', String, 'string'),
            #  Special
            (r'\[\]', Keyword.Type),
            (r'\(\)', Name.Builtin),
            (r'[][(),;`{}]', Punctuation),
        ],
        'import': [
            # Import statements
            (r'\s+', Text),
            (r'"', String, 'string'),
            # after "funclist" state
            (r'\)', Punctuation, '#pop'),
            (r'qualified\b', Keyword),
            # import X as Y
            (r'([A-Z][\w.]*)(\s+)(as)(\s+)([A-Z][\w.]*)',
             bygroups(Name.Namespace, Text, Keyword, Text, Name), '#pop'),
            # import X hiding (functions)
            (r'([A-Z][\w.]*)(\s+)(hiding)(\s+)(\()',
             bygroups(Name.Namespace, Text, Keyword, Text, Punctuation), 'funclist'),
            # import X (functions)
            (r'([A-Z][\w.]*)(\s+)(\()',
             bygroups(Name.Namespace, Text, Punctuation), 'funclist'),
            # import X
            (r'[\w.]+', Name.Namespace, '#pop'),
        ],
        'module': [
            (r'\s+', Text),
            (r'([A-Z][\w.]*)(\s+)(\()',
             bygroups(Name.Namespace, Text, Punctuation), 'funclist'),
            (r'[A-Z][\w.]*', Name.Namespace, '#pop'),
        ],
        'funclist': [
            (r'\s+', Text),
            (r'[A-Z]\w*', Keyword.Type),
            (r'(_[\w\']+|[a-z][\w\']*)', Name.Function),
            # TODO: these don't match the comments in docs, remove.
            #(r'--(?![!#$%&*+./<=>?@^|_~:\\]).*?$', Comment.Single),
            #(r'{-', Comment.Multiline, 'comment'),
            (r',', Punctuation),
            (r'[:!#$%&*+.\\/<=>?@^|~-]+', Operator),
            # (HACK, but it makes sense to push two instances, believe me)
            (r'\(', Punctuation, ('funclist', 'funclist')),
            (r'\)', Punctuation, '#pop:2'),
        ],
        'comment': [
            # Multiline Comments
            (r'[^/*]+', Comment.Multiline),
            (r'/\*', Comment.Multiline, '#push'),
            (r'\*/', Comment.Multiline, '#pop'),
            (r'[*/]', Comment.Multiline),
        ],
        'character': [
            # Allows multi-chars, incorrectly.
            (r"[^\\']'", String.Char, '#pop'),
            (r"\\", String.Escape, 'escape'),
            ("'", String.Char, '#pop'),
        ],
        'string': [
            (r'[^\\"]+', String),
            (r"\\", String.Escape, 'escape'),
            ('"', String, '#pop'),
        ],
        'escape': [
            (r'[abfnrtv"\'&\\]', String.Escape, '#pop'),
            (r'\^[][A-Z@^_]', String.Escape, '#pop'),
            ('|'.join(ascii), String.Escape, '#pop'),
            (r'o[0-7]+', String.Escape, '#pop'),
            (r'x[\da-fA-F]+', String.Escape, '#pop'),
            (r'\d+', String.Escape, '#pop'),
            (r'\s+\\', String.Escape, '#pop'),
        ],
    }

    EXTRA_KEYWORDS = set(('join', 'split', 'reverse', 'transpose', 'width',
                          'length', 'tail', '<<', '>>', '<<<', '>>>', 'const',
                          'reg', 'par', 'seq', 'ASSERT', 'undefined', 'error',
                          'trace'))

    def get_tokens_unprocessed(self, text):
        stack = ['root']
        for index, token, value in \
                RegexLexer.get_tokens_unprocessed(self, text, stack):
            if token is Name and value in self.EXTRA_KEYWORDS:
                yield index, Name.Builtin, value
            else:
                yield index, token, value

class LiterateLexer(Lexer):
    """
    Base class for lexers of literate file formats based on LaTeX or Bird-style
    (prefixing each code line with ">").

    Additional options accepted:

    `litstyle`
        If given, must be ``"bird"`` or ``"latex"``.  If not given, the style
        is autodetected: if the first non-whitespace character in the source
        is a backslash or percent character, LaTeX is assumed, else Bird.
    """

    bird_re = re.compile(r'(>[ \t]*)(.*\n)')

    def __init__(self, baselexer, **options):
        self.baselexer = baselexer
        Lexer.__init__(self, **options)

    def get_tokens_unprocessed(self, text):
        style = self.options.get('litstyle')
        if style is None:
            style = (text.lstrip()[0:1] in '%\\') and 'latex' or 'bird'

        code = ''
        insertions = []
        if style == 'bird':
            # bird-style
            for match in line_re.finditer(text):
                line = match.group()
                m = self.bird_re.match(line)
                if m:
                    insertions.append((len(code),
                                       [(0, Comment.Special, m.group(1))]))
                    code += m.group(2)
                else:
                    insertions.append((len(code), [(0, Text, line)]))
        else:
            # latex-style
            from pygments.lexers.markup import TexLexer
            lxlexer = TexLexer(**self.options)
            codelines = 0
            latex = ''
            for match in line_re.finditer(text):
                line = match.group()
                if codelines:
                    if line.lstrip().startswith('\\end{code}'):
                        codelines = 0
                        latex += line
                    else:
                        code += line
                elif line.lstrip().startswith('\\begin{code}'):
                    codelines = 1
                    latex += line
                    insertions.append((len(code),
                                       list(lxlexer.get_tokens_unprocessed(latex))))
                    latex = ''
                else:
                    latex += line
            insertions.append((len(code),
                               list(lxlexer.get_tokens_unprocessed(latex))))
        for item in do_insertions(insertions, self.baselexer.get_tokens_unprocessed(code)):
            yield item

class LiterateHaskellLexer(LiterateLexer):
    """
    For Literate Haskell (Bird-style or LaTeX) source.

    Additional options accepted:

    `litstyle`
        If given, must be ``"bird"`` or ``"latex"``.  If not given, the style
        is autodetected: if the first non-whitespace character in the source
        is a backslash or percent character, LaTeX is assumed, else Bird.

    .. versionadded:: 0.9
    """
    name = 'Literate Haskell'
    aliases = ['lhs', 'literate-haskell', 'lhaskell']
    filenames = ['*.lhs']
    mimetypes = ['text/x-literate-haskell']

    def __init__(self, **options):
        hslexer = HaskellLexer(**options)
        LiterateLexer.__init__(self, hslexer, **options)

class LiterateIdrisLexer(LiterateLexer):
    """
    For Literate Idris (Bird-style or LaTeX) source.

    Additional options accepted:

    `litstyle`
        If given, must be ``"bird"`` or ``"latex"``.  If not given, the style
        is autodetected: if the first non-whitespace character in the source
        is a backslash or percent character, LaTeX is assumed, else Bird.

    .. versionadded:: 2.0
    """
    name = 'Literate Idris'
    aliases = ['lidr', 'literate-idris', 'lidris']
    filenames = ['*.lidr']
    mimetypes = ['text/x-literate-idris']

    def __init__(self, **options):
        hslexer = IdrisLexer(**options)
        LiterateLexer.__init__(self, hslexer, **options)

class LiterateAgdaLexer(LiterateLexer):
    """
    For Literate Agda source.

    Additional options accepted:

    `litstyle`
        If given, must be ``"bird"`` or ``"latex"``.  If not given, the style
        is autodetected: if the first non-whitespace character in the source
        is a backslash or percent character, LaTeX is assumed, else Bird.

    .. versionadded:: 2.0
    """
    name = 'Literate Agda'
    aliases = ['lagda', 'literate-agda']
    filenames = ['*.lagda']
    mimetypes = ['text/x-literate-agda']

    def __init__(self, **options):
        agdalexer = AgdaLexer(**options)
        LiterateLexer.__init__(self, agdalexer, litstyle='latex', **options)

class LiterateCryptolLexer(LiterateLexer):
    """
    For Literate Cryptol (Bird-style or LaTeX) source.

    Additional options accepted:

    `litstyle`
        If given, must be ``"bird"`` or ``"latex"``.  If not given, the style
        is autodetected: if the first non-whitespace character in the source
        is a backslash or percent character, LaTeX is assumed, else Bird.

    .. versionadded:: 2.0
    """
    name = 'Literate Cryptol'
    aliases = ['lcry', 'literate-cryptol', 'lcryptol']
    filenames = ['*.lcry']
    mimetypes = ['text/x-literate-cryptol']

    def __init__(self, **options):
        crylexer = CryptolLexer(**options)
        LiterateLexer.__init__(self, crylexer, **options)

class KokaLexer(RegexLexer):
    """
    Lexer for the `Koka <http://koka.codeplex.com>`_
    language.

    .. versionadded:: 1.6
    """

    name = 'Koka'
    aliases = ['koka']
    filenames = ['*.kk', '*.kki']
    mimetypes = ['text/x-koka']

    keywords = [
        'infix', 'infixr', 'infixl',
        'type', 'cotype', 'rectype', 'alias',
        'struct', 'con',
        'fun', 'function', 'val', 'var',
        'external',
        'if', 'then', 'else', 'elif', 'return', 'match',
        'private', 'public', 'private',
        'module', 'import', 'as',
        'include', 'inline',
        'rec',
        'try', 'yield', 'enum',
        'interface', 'instance',
    ]

    # keywords that are followed by a type
    typeStartKeywords = [
        'type', 'cotype', 'rectype', 'alias', 'struct', 'enum',
    ]

    # keywords valid in a type
    typekeywords = [
        'forall', 'exists', 'some', 'with',
    ]

    # builtin names and special names
    builtin = [
        'for', 'while', 'repeat',
        'foreach', 'foreach-indexed',
        'error', 'catch', 'finally',
        'cs', 'js', 'file', 'ref', 'assigned',
    ]

    # symbols that can be in an operator
    symbols = r'[$%&*+@!/\\^~=.:\-?|<>]+'

    # symbol boundary: an operator keyword should not be followed by any of these
    sboundary = '(?!'+symbols+')'

    # name boundary: a keyword should not be followed by any of these
    boundary = '(?![\w/])'

    # koka token abstractions
    tokenType = Name.Attribute
    tokenTypeDef = Name.Class
    tokenConstructor = Generic.Emph

    # main lexer
    tokens = {
        'root': [
            include('whitespace'),

            # go into type mode
            (r'::?' + sboundary, tokenType, 'type'),
            (r'(alias)(\s+)([a-z]\w*)?', bygroups(Keyword, Text, tokenTypeDef),
             'alias-type'),
            (r'(struct)(\s+)([a-z]\w*)?', bygroups(Keyword, Text, tokenTypeDef),
             'struct-type'),
            ((r'(%s)' % '|'.join(typeStartKeywords)) +
             r'(\s+)([a-z]\w*)?', bygroups(Keyword, Text, tokenTypeDef),
             'type'),

            # special sequences of tokens (we use ?: for non-capturing group as
            # required by 'bygroups')
            (r'(module)(\s+)(interface\s+)?((?:[a-z]\w*/)*[a-z]\w*)',
             bygroups(Keyword, Text, Keyword, Name.Namespace)),
            (r'(import)(\s+)((?:[a-z]\w*/)*[a-z]\w*)'
             r'(?:(\s*)(=)(\s*)((?:qualified\s*)?)'
             r'((?:[a-z]\w*/)*[a-z]\w*))?',
             bygroups(Keyword, Text, Name.Namespace, Text, Keyword, Text,
                      Keyword, Name.Namespace)),

            (r'(^(?:(?:public|private)\s*)?(?:function|fun|val))'
             r'(\s+)([a-z]\w*|\((?:' + symbols + r'|/)\))',
             bygroups(Keyword, Text, Name.Function)),
            (r'(^(?:(?:public|private)\s*)?external)(\s+)(inline\s+)?'
             r'([a-z]\w*|\((?:' + symbols + r'|/)\))',
             bygroups(Keyword, Text, Keyword, Name.Function)),

            # keywords
            (r'(%s)' % '|'.join(typekeywords) + boundary, Keyword.Type),
            (r'(%s)' % '|'.join(keywords) + boundary, Keyword),
            (r'(%s)' % '|'.join(builtin) + boundary, Keyword.Pseudo),
            (r'::?|:=|\->|[=.]' + sboundary, Keyword),

            # names
            (r'((?:[a-z]\w*/)*)([A-Z]\w*)',
             bygroups(Name.Namespace, tokenConstructor)),
            (r'((?:[a-z]\w*/)*)([a-z]\w*)', bygroups(Name.Namespace, Name)),
            (r'((?:[a-z]\w*/)*)(\((?:' + symbols + r'|/)\))',
             bygroups(Name.Namespace, Name)),
            (r'_\w*', Name.Variable),

            # literal string
            (r'@"', String.Double, 'litstring'),

            # operators
            (symbols + "|/(?![*/])", Operator),
            (r'`', Operator),
            (r'[{}()\[\];,]', Punctuation),

            # literals. No check for literal characters with len > 1
            (r'[0-9]+\.[0-9]+([eE][\-+]?[0-9]+)?', Number.Float),
            (r'0[xX][0-9a-fA-F]+', Number.Hex),
            (r'[0-9]+', Number.Integer),

            (r"'", String.Char, 'char'),
            (r'"', String.Double, 'string'),
        ],

        # type started by alias
        'alias-type': [
            (r'=', Keyword),
            include('type')
        ],

        # type started by struct
        'struct-type': [
            (r'(?=\((?!,*\)))', Punctuation, '#pop'),
            include('type')
        ],

        # type started by colon
        'type': [
            (r'[(\[<]', tokenType, 'type-nested'),
            include('type-content')
        ],

        # type nested in brackets: can contain parameters, comma etc.
        'type-nested': [
            (r'[)\]>]', tokenType, '#pop'),
            (r'[(\[<]', tokenType, 'type-nested'),
            (r',', tokenType),
            (r'([a-z]\w*)(\s*)(:)(?!:)',
             bygroups(Name, Text, tokenType)),  # parameter name
            include('type-content')
        ],

        # shared contents of a type
        'type-content': [
            include('whitespace'),

            # keywords
            (r'(%s)' % '|'.join(typekeywords) + boundary, Keyword),
            (r'(?=((%s)' % '|'.join(keywords) + boundary + '))',
             Keyword, '#pop'),  # need to match because names overlap...

            # kinds
            (r'[EPHVX]' + boundary, tokenType),

            # type names
            (r'[a-z][0-9]*(?![\w/])', tokenType),
            (r'_\w*', tokenType.Variable),  # Generic.Emph
            (r'((?:[a-z]\w*/)*)([A-Z]\w*)',
             bygroups(Name.Namespace, tokenType)),
            (r'((?:[a-z]\w*/)*)([a-z]\w+)',
             bygroups(Name.Namespace, tokenType)),

            # type keyword operators
            (r'::|->|[.:|]', tokenType),

            # catchall
            default('#pop')
        ],

        # comments and literals
        'whitespace': [
            (r'\n\s*#.*$', Comment.Preproc),
            (r'\s+', Text),
            (r'/\*', Comment.Multiline, 'comment'),
            (r'//.*$', Comment.Single)
        ],
        'comment': [
            (r'[^/*]+', Comment.Multiline),
            (r'/\*', Comment.Multiline, '#push'),
            (r'\*/', Comment.Multiline, '#pop'),
            (r'[*/]', Comment.Multiline),
        ],
        'litstring': [
            (r'[^"]+', String.Double),
            (r'""', String.Escape),
            (r'"', String.Double, '#pop'),
        ],
        'string': [
            (r'[^\\"\n]+', String.Double),
            include('escape-sequence'),
            (r'["\n]', String.Double, '#pop'),
        ],
        'char': [
            (r'[^\\\'\n]+', String.Char),
            include('escape-sequence'),
            (r'[\'\n]', String.Char, '#pop'),
        ],
        'escape-sequence': [
            (r'\\[nrt\\"\']', String.Escape),
            (r'\\x[0-9a-fA-F]{2}', String.Escape),
            (r'\\u[0-9a-fA-F]{4}', String.Escape),
            # Yes, \U literals are 6 hex digits.
            (r'\\U[0-9a-fA-F]{6}', String.Escape)
        ]
    }


import typing

import gym

on_state_change_type = typing.Callable[[
    gym.Space,  # State
    gym.Space,  # Action
    float,  # Reward
    gym.Space,  # New state
    bool,  # Is done
    typing.Optional[object],  # Info
], type(None)]

def _noop_function(*args, **kwargs):
    pass

class BasePolicy(object):
    """A class that represents an exploration policy for Go-Explore"""

    def __init__(self, environment: gym.Env):
        """Create exploration policy.

        :param environment: OpenAI Gym environment that should be explored
        """
        self.environment = environment
        self._on_action = _noop_function

    @property
    def on_action(self):
        return self._on_action

    @on_action.setter
    def on_action(self, new_on_action: on_state_change_type):
        self._on_action = new_on_action

    def _environment_act(self, current_state: gym.Space, action: gym.Space):
        result = self.environment.step(action)
        new_state, reward, done, info = result
        self._on_action(current_state, action, reward, new_state, done, info)
        return result

    def explore(self, current_state: gym.Space):
        """Explore from current state.

        This method should explore from current_state using the exploration
        policy. It can be e.g. random actions, exploration through curiosity,
        etc.

        The environment should be in current_state in order for this method to
        work properly.

        :param current_state: Current state of the environment
        :returns Latest tuple from env.step call (or None if not explored)
        """
        raise NotImplementedError


from disco import util
from discodb import DiscoDB, Q
from disco.worker.task_io import task_output_stream

def Open(url, task=None):
    if task:
        disco_data = task.disco_data
        ddfs_data = task.ddfs_data
    else:
        from disco.settings import DiscoSettings
        settings = DiscoSettings()
        disco_data = settings['DISCO_DATA']
        ddfs_data = settings['DDFS_DATA']
    scheme, netloc, rest = util.urlsplit(url)
    path, rest = rest.split('!', 1) if '!' in rest else (rest, '')
    discodb = DiscoDB.load(open(util.localize(path, disco_data=disco_data,
                                ddfs_data=ddfs_data)))

    if rest:
        method_name, arg = rest.split('/', 1) if '/' in rest else (rest, None)
        method = getattr(discodb, method_name)
        if method_name in ('metaquery', 'query'):
            return method(Q.urlscan(arg))
        return method(*filter(None, arg))
    return discodb

def input_stream(fd, size, url, params):
    return Open(url, task=globals().get('Task')), size, url

class DiscoDBOutput(object):
    def __init__(self, stream, params):
        from discodb import DiscoDBConstructor
        self.discodb_constructor = DiscoDBConstructor()
        self.stream = stream
        self.params = params
        self.path = stream.path

    def add(self, key, val):
        self.discodb_constructor.add(key, val)

    def close(self):
        def flags():
            return dict((flag, getattr(self.params, flag))
                        for flag in ('unique_items', 'disable_compression')
                        if hasattr(self.params, flag))
        self.discodb_constructor.finalize(**flags()).dump(self.stream)

def discodb_output(stream, partition, url, params):
    return DiscoDBOutput(stream, params), 'discodb:{0}'.format(url.split(':', 1)[1])

discodb_stream = (task_output_stream, discodb_output)


"""
Python interface module for OSQP solver v0.6.2.post5
"""
from __future__ import print_function
from builtins import object
import osqp._osqp as _osqp  # Internal low level module
import numpy as np
import scipy.sparse as spa
from warnings import warn
from platform import system
import osqp.codegen as cg
import osqp.utils as utils
import sys
import qdldl

class OSQP(object):
    def __init__(self):
        self._model = _osqp.OSQP()

    def version(self):
        return self._model.version()

    def setup(self, P=None, q=None, A=None, l=None, u=None, **settings):
        """
        Setup OSQP solver problem of the form

        minimize     1/2 x' * P * x + q' * x
        subject to   l <= A * x <= u

        solver settings can be specified as additional keyword arguments
        """
        # TODO(bart): this will be unnecessary when the derivative will be in C
        self._derivative_cache = {'P': P, 'q': q, 'A': A, 'l': l, 'u': u}

        unpacked_data, settings = utils.prepare_data(P, q, A, l, u, **settings)
        self._model.setup(*unpacked_data, **settings)

    def update(self, q=None, l=None, u=None,
               Px=None, Px_idx=np.array([]), Ax=None, Ax_idx=np.array([])):
        """
        Update OSQP problem arguments
        """

        # get problem dimensions
        (n, m) = self._model.dimensions()

        # check consistency of the input arguments
        if q is not None and len(q) != n:
            raise ValueError("q must have length n")
        if l is not None:
            if not isinstance(l, np.ndarray):
                raise TypeError("l must be numpy.ndarray, not %s" %
                                type(l).__name__)
            elif len(l) != m:
                raise ValueError("l must have length m")
            # Convert values to -OSQP_INFTY
            l = np.maximum(l, -_osqp.constant('OSQP_INFTY'))
        if u is not None:
            if not isinstance(u, np.ndarray):
                raise TypeError("u must be numpy.ndarray, not %s" %
                                type(u).__name__)
            elif len(u) != m:
                raise ValueError("u must have length m")
            # Convert values to OSQP_INFTY
            u = np.minimum(u, _osqp.constant('OSQP_INFTY'))
        if Ax is None:
            if len(Ax_idx) > 0:
                raise ValueError("Vector Ax has not been specified")
        else:
            if len(Ax_idx) > 0 and len(Ax) != len(Ax_idx):
                raise ValueError("Ax and Ax_idx must have the same lengths")
        if Px is None:
            if len(Px_idx) > 0:
                raise ValueError("Vector Px has not been specified")
        else:
            if len(Px_idx) > 0 and len(Px) != len(Px_idx):
                raise ValueError("Px and Px_idx must have the same lengths")
        if q is None and l is None and u is None and Px is None and Ax is None:
            raise ValueError("No updatable data has been specified")

        # update linear cost
        if q is not None:
            self._model.update_lin_cost(q)

        # update lower bound
        if l is not None and u is None:
            self._model.update_lower_bound(l)

        # update upper bound
        if u is not None and l is None:
            self._model.update_upper_bound(u)

        # update bounds
        if l is not None and u is not None:
            self._model.update_bounds(l, u)

        # update matrix P
        if Px is not None and Ax is None:
            self._model.update_P(Px, Px_idx, len(Px))

        # update matrix A
        if Ax is not None and Px is None:
            self._model.update_A(Ax, Ax_idx, len(Ax))

        # update matrices P and A
        if Px is not None and Ax is not None:
            self._model.update_P_A(Px, Px_idx, len(Px), Ax, Ax_idx, len(Ax))

        # TODO(bart): this will be unnecessary when the derivative will be in C
        # update problem data in self._derivative_cache
        if q is not None:
            self._derivative_cache["q"] = q

        if l is not None:
            self._derivative_cache["l"] = l

        if u is not None:
            self._derivative_cache["u"] = u

        if Px is not None:
            if Px_idx.size == 0:
                self._derivative_cache["P"].data = Px
            else:
                self._derivative_cache["P"].data[Px_idx] = Px

        if Ax is not None:
            if Ax_idx.size == 0:
                self._derivative_cache["A"].data = Ax
            else:
                self._derivative_cache["A"].data[Ax_idx] = Ax

        # delete results from self._derivative_cache to prohibit
        # taking the derivative of unsolved problems
        if "results" in self._derivative_cache.keys():
            del self._derivative_cache["results"]

    def update_settings(self, **kwargs):
        """
        Update OSQP solver settings

        It is possible to change: 'max_iter', 'eps_abs', 'eps_rel',
                                  'eps_prim_inf', 'eps_dual_inf', 'rho'
                                  'alpha', 'delta', 'polish',
                                  'polish_refine_iter',
                                  'verbose', 'scaled_termination',
                                  'check_termination', 'time_limit',
        """

        # get arguments
        max_iter = kwargs.pop('max_iter', None)
        eps_abs = kwargs.pop('eps_abs', None)
        eps_rel = kwargs.pop('eps_rel', None)
        eps_prim_inf = kwargs.pop('eps_prim_inf', None)
        eps_dual_inf = kwargs.pop('eps_dual_inf', None)
        rho = kwargs.pop('rho', None)
        alpha = kwargs.pop('alpha', None)
        delta = kwargs.pop('delta', None)
        polish = kwargs.pop('polish', None)
        polish_refine_iter = kwargs.pop('polish_refine_iter', None)
        verbose = kwargs.pop('verbose', None)
        scaled_termination = kwargs.pop('scaled_termination', None)
        check_termination = kwargs.pop('check_termination', None)
        warm_start = kwargs.pop('warm_start', None)
        time_limit = kwargs.pop('time_limit', None)

        # update them
        if max_iter is not None:
            self._model.update_max_iter(max_iter)

        if eps_abs is not None:
            self._model.update_eps_abs(eps_abs)

        if eps_rel is not None:
            self._model.update_eps_rel(eps_rel)

        if eps_prim_inf is not None:
            self._model.update_eps_prim_inf(eps_prim_inf)

        if eps_dual_inf is not None:
            self._model.update_eps_dual_inf(eps_dual_inf)

        if rho is not None:
            self._model.update_rho(rho)

        if alpha is not None:
            self._model.update_alpha(alpha)

        if delta is not None:
            self._model.update_delta(delta)

        if polish is not None:
            self._model.update_polish(polish)

        if polish_refine_iter is not None:
            self._model.update_polish_refine_iter(polish_refine_iter)

        if verbose is not None:
            self._model.update_verbose(verbose)

        if scaled_termination is not None:
            self._model.update_scaled_termination(scaled_termination)

        if check_termination is not None:
            self._model.update_check_termination(check_termination)

        if warm_start is not None:
            self._model.update_warm_start(warm_start)

        if time_limit is not None:
            self._model.update_time_limit(time_limit)

        if max_iter is None and \
           eps_abs is None and \
           eps_rel is None and \
           eps_prim_inf is None and \
           eps_dual_inf is None and \
           rho is None and \
           alpha is None and \
           delta is None and \
           polish is None and \
           polish_refine_iter is None and \
           verbose is None and \
           scaled_termination is None and \
           check_termination is None and \
           warm_start is None:
            raise ValueError("No updatable settings has been specified!")

    def solve(self):
        """
        Solve QP Problem
        """
        # Solve QP
        results = self._model.solve()

        # TODO(bart): this will be unnecessary when the derivative will be in C
        self._derivative_cache['results'] = results

        return results

    def warm_start(self, x=None, y=None):
        """
        Warm start primal or dual variables
        """
        # get problem dimensions
        (n, m) = self._model.dimensions()

        if x is not None:
            if len(x) != n:
                raise ValueError("Wrong dimension for variable x")

            if y is None:
                self._model.warm_start_x(x)

        if y is not None:
            if len(y) != m:
                raise ValueError("Wrong dimension for variable y")

            if x is None:
                self._model.warm_start_y(y)

        if x is not None and y is not None:
            self._model.warm_start(x, y)

        if x is None and y is None:
            raise ValueError("Unrecognized fields")

    def codegen(self, folder, project_type='', parameters='vectors',
                python_ext_name='emosqp', force_rewrite=False, compile_python_ext=True,
                FLOAT=False, LONG=True):
        """
        Generate embeddable C code for the problem
        """

        # Check parameters arguments
        if parameters == 'vectors':
            embedded = 1
        elif parameters == 'matrices':
            embedded = 2
        else:
            raise ValueError("Unknown value of 'parameters' argument.")

        # Set float and long flags
        if FLOAT:
            float_flag = 'ON'
        else:
            float_flag = 'OFF'
        if LONG:
            long_flag = 'ON'
        else:
            long_flag = 'OFF'

        # Check project_type argument
        expectedProject = ('', 'Makefile', 'MinGW Makefiles',
                           'Unix Makefiles', 'CodeBlocks', 'Xcode')
        if project_type not in expectedProject:
            raise ValueError("Unknown value of 'project_type' argument.")

        if project_type == 'Makefile':
            if system() == 'Windows':
                project_type = 'MinGW Makefiles'
            elif system() == 'Linux' or system() == 'Darwin':
                project_type = 'Unix Makefiles'

        # Convert workspace to Python
        sys.stdout.write("Getting workspace from OSQP object... \t\t\t\t")
        sys.stdout.flush()
        work = self._model._get_workspace()
        print("[done]")

        # Generate code with codegen module
        cg.codegen(work, folder, python_ext_name, project_type, compile_python_ext,
                   embedded, force_rewrite, float_flag, long_flag)

    def derivative_iterative_refinement(self, rhs, max_iter=20, tol=1e-12):
        M = self._derivative_cache['M']

        # Prefactor
        solver = self._derivative_cache['solver']

        sol = solver.solve(rhs)
        for k in range(max_iter):
            delta_sol = solver.solve(rhs - M @ sol)
            sol = sol + delta_sol

            if np.linalg.norm(M @ sol - rhs) < tol:
                break

        if k == max_iter - 1:
            warn("max_iter iterative refinement reached.")

        return sol

    def adjoint_derivative(self, dx=None, dy_u=None, dy_l=None,
                           P_idx=None, A_idx=None, eps_iter_ref=1e-04):
        """
        Compute adjoint derivative after solve.
        """

        P, q = self._derivative_cache['P'], self._derivative_cache['q']
        A = self._derivative_cache['A']
        l, u = self._derivative_cache['l'], self._derivative_cache['u']

        try:
            results = self._derivative_cache['results']
        except KeyError:
            raise ValueError("Problem has not been solved. "
                             "You cannot take derivatives. "
                             "Please call the solve function.")

        if results.info.status != "solved":
            raise ValueError("Problem has not been solved to optimality. "
                             "You cannot take derivatives")

        m, n = A.shape
        x = results.x
        y = results.y
        y_u = np.maximum(y, 0)
        y_l = -np.minimum(y, 0)

        if A_idx is None:
            A_idx = A.nonzero()

        if P_idx is None:
            P_idx = P.nonzero()

        if dy_u is None:
            dy_u = np.zeros(m)
        if dy_l is None:
            dy_l = np.zeros(m)

        # Make sure M matrix exists
        if 'M' not in self._derivative_cache:
            # Multiply second-third row by diag(y_u)^-1 and diag(y_l)^-1
            # to make the matrix symmetric
            inv_dia_y_u = spa.diags(np.reciprocal(y_u + 1e-20))
            inv_dia_y_l = spa.diags(np.reciprocal(y_l + 1e-20))
            M = spa.bmat([
                [P,            A.T,                  -A.T],
                [A, spa.diags(A @ x - u) @ inv_dia_y_u, None],
                [-A, None, spa.diags(l - A @ x) @ inv_dia_y_l]
            ], format='csc')
            delta = spa.bmat([[eps_iter_ref * spa.eye(n), None],
                              [None, -eps_iter_ref * spa.eye(2 * m)]],
                             format='csc')
            self._derivative_cache['M'] = M
            self._derivative_cache['solver'] = qdldl.Solver(M + delta)

        rhs = - np.concatenate([dx, dy_u, dy_l])

        r_sol = self.derivative_iterative_refinement(rhs)

        r_x, r_yu, r_yl = np.split(r_sol, [n, n+m])

        # Extract derivatives for the constraints
        rows, cols = A_idx
        dA_vals = (y_u[rows] - y_l[rows]) * r_x[cols] + \
            (r_yu[rows] - r_yl[rows]) * x[cols]
        dA = spa.csc_matrix((dA_vals, (rows, cols)), shape=A.shape)
        du = - r_yu
        dl = r_yl

        # Extract derivatives for the cost (P, q)
        rows, cols = P_idx
        dP_vals = .5 * (r_x[rows] * x[cols] + r_x[cols] * x[rows])
        dP = spa.csc_matrix((dP_vals, P_idx), shape=P.shape)
        dq = r_x

        return (dP, dq, dA, dl, du)


"""
Base interface for a reader class
"""
import numpy as np
import logging
from pathlib import Path
import matplotlib.pyplot as plt
import pandas as pd

logger = logging.getLogger(__name__)
import PyFLOTRAN.utils.SubFishModule as subfish
import seaborn as sns

class BaseReader:
    data: np.ndarray  # Hint of self.data array
    info: dict

    def __init__(self, filename=None, header=False, **kwargs):
        self.filename = Path(filename)
        self.info = {"reader": {}}
        self.data = None
        self.header = header
        self.__dict__.update(kwargs)
        self.open_file(filename)

    def read_file(self, opened_file):
        """
        Reads the data and stores it inside the class
        :return:
        """
        pass

    def open_file(self, filename):
        with open(filename) as opened_file:
            if self.header:
                opened_file.readline()  # For now, skips the header if it has
            self.read_file(opened_file)
        self.build_info()

    def read_header(self, opened_file):
        """
        Reads the header of the file
        :return:
        """
        pass

    def get_data(self) -> np.ndarray:
        """
        Outputs the read data
        :return:
        """
        return np.array(0)

    def build_info(self):
        """
        Generates a dictionary containing the basic info of the read data
        :return:
        """
        self.info = {}

    def global_coords_to_local(self, x_local_to_global, y_local_to_global):
        """Converts global data coordinates into local"""
        assert len(self.data.shape) >= 2 and self.data.shape[1] >= 2, "Error in data shape"
        self.data[:, 0] -= x_local_to_global
        self.data[:, 1] -= y_local_to_global

    def local_coords_to_global(self, x_local_to_global, y_local_to_global):
        """Converts local data coordinates into global"""
        assert len(self.data.shape) >= 2 and self.data.shape[1] >= 2, "Error in data shape"
        self.data[:, 0] += x_local_to_global
        self.data[:, 1] += y_local_to_global

    def dump_to_csv(self, output_file, delimiter=","):
        """
        Writes the data into a csv file
        :param output_file:
        :return:
        """
        print(f"Starting dump into {output_file}")
        np.savetxt(output_file, self.get_data(), delimiter=delimiter)
        print(f"The data has been properly exported to the {output_file} file")

    def create_postprocess_dict(self):
        self.postprocessing_dict = Path().cwd() / "postprocess"
        self.postprocessing_dict.mkdir(exist_ok=True)

    def generate_subfish_data(self, subfish_dict: dict, unit_factor=1/(365 * 24), unit_name='d') -> pd.DataFrame:
        """
        This method reads a given subfish dict and returns the calculated results
        Args:
            subfish_dict: dictionary containing subfish module parameters
            unit_factor: factor multiplyin the results (originally computed in seconds)
        Returns:
            Array containing the times and computed values
        """
        if not subfish_dict:
            raise AttributeError('Please, provide a suitable subfish dict object')
        tang_sol = subfish.calculate_tang(subfish_dict)
        tang_sol[0] *= unit_factor
        tang_sol_pd = pd.DataFrame({f'Time [{unit_name}]': tang_sol[0],
                                    'Result [M]': tang_sol[1]
                                    }
                                   )
        return tang_sol_pd


#!/usr/bin/env python
from __future__ import with_statement, division
from Tkinter import Tk, Frame, Button, Canvas, Label, NW
from array import array
from PIL import Image as PILImage, ImageTk, ImageDraw

def tileFromPlanar(s, opaqueBase=0):
    nPlanes = len(s) // 8
    im = PILImage.new('P', (8, 8))
    pixels = im.load()
    if pixels is None:
        print "Ouch!", repr(im)
    for y in range(8):
        planedata = [ord(c) for c in s[y::8]]
        for x in range(8):
            c = 0
            for i in range(nPlanes):
                if planedata[i] & (0x80 >> x):
                    c |= 1 << i
            pixels[x, y] = c + opaqueBase if c > 0 else c
    return im

def renderAttrPicker(tileImg, palette, value):
    im = PILImage.new('P', (128, 32))
    im.putpalette(palette)
    for i in range(4):
        x = i * 32
        t = 0 if i == value else 3
        im.paste(i * 4 + 1, (x + 16, t, x + 32 - t, 16))
        im.paste(i * 4 + 2, (x + t, 16, x + 16, 32 - t))
        im.paste(i * 4 + 3, (x + 16, 16, x + 32 - t, 32 - t))
    im.paste(tileImg.resize((16, 16)), (value * 32, 0))
    #im.show()
    return im

def renderChrFile(tiledata, palette, opaqueBase=0):
    tiles = [tiledata[i:i + 16] for i in range(0, len(tiledata), 16)]
    rows = [tiles[i:i + 16] for i in range(0, len(tiles), 16)]
    h = len(rows) * 8
    w = 128

    im = PILImage.new('P', (w, h))
    y = 0
    #palette = [0, 0, 0, 153, 102, 0, 102, 204, 0, 255, 204, 0]
    im.putpalette(palette)
    for row in rows:
        x = 0
        for tiledata in row:
            tile = tileFromPlanar(tiledata, opaqueBase)
            im.paste(tile, (x, y))
            x += 8
        y += 8
    return im

class NamFile:
    defaultNESPalette = \
        "\x0f\x00\x10\x30\x0f\x12\x1a\x30\x0f\x1a\x2c\x30\x0f\x12\x14\x30"
    nesclut = [
        (0x80,0x80,0x80), (0x00,0x00,0xBB), (0x37,0x00,0xBF), (0x84,0x00,0xA6),
        (0xBB,0x00,0x6A), (0xB7,0x00,0x1E), (0xB3,0x00,0x00), (0x91,0x26,0x00),
        (0x7B,0x2B,0x00), (0x00,0x3E,0x00), (0x00,0x48,0x0D), (0x00,0x3C,0x22),
        (0x00,0x2F,0x66), (0x00,0x00,0x00), (0x05,0x05,0x05), (0x05,0x05,0x05),

        (0xC8,0xC8,0xC8), (0x00,0x59,0xFF), (0x44,0x3C,0xFF), (0xB7,0x33,0xCC),
        (0xFF,0x33,0xAA), (0xFF,0x37,0x5E), (0xFF,0x37,0x1A), (0xD5,0x4B,0x00),
        (0xC4,0x62,0x00), (0x3C,0x7B,0x00), (0x1E,0x84,0x15), (0x00,0x95,0x66),
        (0x00,0x84,0xC4), (0x11,0x11,0x11), (0x09,0x09,0x09), (0x09,0x09,0x09),

        (0xFF,0xFF,0xFF), (0x00,0x95,0xFF), (0x6F,0x84,0xFF), (0xD5,0x6F,0xFF),
        (0xFF,0x77,0xCC), (0xFF,0x6F,0x99), (0xFF,0x7B,0x59), (0xFF,0x91,0x5F),
        (0xFF,0xA2,0x33), (0xA6,0xBF,0x00), (0x51,0xD9,0x6A), (0x4D,0xD5,0xAE),
        (0x00,0xD9,0xFF), (0x66,0x66,0x66), (0x0D,0x0D,0x0D), (0x0D,0x0D,0x0D),

        (0xFF,0xFF,0xFF), (0x84,0xBF,0xFF), (0xBB,0xBB,0xFF), (0xD0,0xBB,0xFF),
        (0xFF,0xBF,0xEA), (0xFF,0xBF,0xCC), (0xFF,0xC4,0xB7), (0xFF,0xCC,0xAE),
        (0xFF,0xD9,0xA2), (0xCC,0xE1,0x99), (0xAE,0xEE,0xB7), (0xAA,0xF7,0xEE),
        (0xB3,0xEE,0xFF), (0xDD,0xDD,0xDD), (0x11,0x11,0x11), (0x11,0x11,0x11)
    ]

    def __init__(self, chrfilename, palfilename='', namfilename=None):
        self.loadchr(chrfilename)
        self.loadpal(palfilename)
        self.loadnam(namfilename)

    def loadchr(self, chrfilename):
        with open(chrfilename, 'rb') as infp:
            chrdata = infp.read(4096)
            self.chrdata = chrdata
        if len(self.chrdata) != 4096:
            raise ValueError("not enough data for pattern table")

    def loadpal(self, palfilename):
        self.palfilename = palfilename
        try:
            with open(palfilename, 'rb') as infp:
                pal = infp.read(16)
            if len(paldata) != 16:
                raise ValueError("not enough data for palette")
        except IOError, e:
            import errno
            if e.errno in (errno.ENOENT, errno.EINVAL):
                pal = self.defaultNESPalette
            else:
                raise
        self.clut = []
        for c in pal:
            self.clut.extend(self.nesclut[ord(c) & 0x3F])

    def loadnam(self, namfilename):
        print "namfilename is", namfilename
        if namfilename is not None:
            try:
                with open(namfilename, 'rb') as infp:
                    namdata = infp.read(1152)
                if ((len(namdata) != 1024
                     and namdata.startswith('\x04\x00'))
                    or namfilename.lower().endswith('.pkb')):
                    print "unpacking"
                    namdata = UnPackBits(namdata[2:]).flush().tostring()
                if len(namdata) != 1024:
                    raise ValueError("not enough data for nametable")
                self.namdata = array('B', namdata)
            except IOError, e:
                import errno
                if e.errno == errno.ENOENT:
                    namfilename = None
                else:
                    raise
        if namfilename is None:
            self.namdata = array('B', [0 for i in range(1024)])
        self.namfilename = namfilename
        self.setUnsaved(False)

    def setUnsaved(self, isSaved):
        import datetime
        self.unsaved = (datetime.datetime.now()
                        if isSaved
                        else False)

    def savenam(self, namfilename=None):
        if namfilename is None:
            namfilename = self.namfilename
        s = self.namdata.tostring()
        if namfilename.lower().endswith('.pkb'):
            s = "\x04\x00" + PackBits(s).flush().tostring()
        with open(namfilename, 'wb') as outfp:
            outfp.write(s)
        self.namfilename = namfilename
        self.setUnsaved(False)

    def getTile(self, x, y):
        if x < 0 or x >= 32 or y < 0 or y >= 30:
            return None
        nameIdx = y * 32 + x
        tileNo = self.namdata[nameIdx]
        attrIdx = (y >> 2) * 8 + (x >> 2) + 960
        attrShift = ((y & 0x02) << 1) | (x & 0x02)
        attrNo = (self.namdata[attrIdx] >> attrShift) & 0x03
        return (tileNo, attrNo)

    def setTile(self, x, y, tileNo, attrNo=None):
        if x < 0 or x >= 32 or y < 0 or y >= 30:
            return
        nameIdx = y * 32 + x
        self.namdata[nameIdx] = tileNo
        if attrNo is not None:
            attrIdx = (y >> 2) * 8 + (x >> 2) + 960
            attrShift = ((y & 0x02) << 1) | (x & 0x02)
            attrByte = (attrNo & 0x03) << attrShift
            attrByte |= self.namdata[attrIdx] & ~(0x03 << attrShift)
            self.namdata[attrIdx] = attrByte

    def getTileData(self, tileNo):
        return self.chrdata[tileNo * 16:tileNo * 16 + 16]

    def renderTile(self, tileNo, attrNo):
        return tileFromPlanar(self.getTileData(tileNo), attrNo * 4)

def build_menubar(parent, mbardata):
    from Tkinter import Menu
    menubar = Menu(parent)
    parent.config(menu=menubar)
    menus = []
    for (label, items) in mbardata:
        menu = Menu(menubar)
        menus.append(menu)
        menubar.add_cascade(label=label, menu=menu)
        for item in items:
            if item == '-':
                menu.add_separator()
            else:
                label = item[0]
                underline = label.find('&')
                if underline >= 0:
                    label = label[:underline] + label[underline+1:]
                else:
                    underline = None
                accelerator = item[2] if len(item) > 2 else None
                menu.add_command(label=label, command=item[1],
                                 accelerator=accelerator,
                                 underline=underline)
    return (menubar, menus)

class TilePicker(Frame):
    def __init__(self, parent, doc, **kw):
        apply(Frame.__init__, (self, parent), kw)
        self.doc = doc
        self.tilePicker = None
        self.attrPicker = None
        self.status = None
        self.curTile = 0
        self.setAttribute(0)
        self.tilePicker = Label(self, image=self.tilePickerPI,
                                width=128, borderwidth=0)
        self.tilePicker.grid(row=0,column=0)
        self.tilePicker.bind("<Button-1>", self.tilePickerCallback)
        self.attrPicker = Label(self, image=self.attrPickerPI,
                                borderwidth=0)
        self.attrPicker.grid(row=1,column=0)
        self.attrPicker.bind("<Button-1>", self.attrPickerCallback)
        self.status = Label(self)
        self.status.grid(row=2,column=0)
        self.setStatus()

    def setAttribute(self, value):
        self.curAttribute = value & 0x03
        self.updateWidgets()

    def updateWidgets(self):
        self.tilePickerImage = renderChrFile(self.doc.chrdata,
                                             self.doc.clut,
                                             self.curAttribute * 4)
        self.tilePickerPI = ImageTk.PhotoImage(self.tilePickerImage)
        if self.tilePicker is not None:
            self.tilePicker.configure(image=self.tilePickerPI)

        previewTile = self.doc.renderTile(self.curTile, self.curAttribute)
        self.attrPickerImage = renderAttrPicker(previewTile,
                                                self.doc.clut,
                                                self.curAttribute)
        self.attrPickerPI = ImageTk.PhotoImage(self.attrPickerImage)
        if self.attrPicker is not None:
            self.attrPicker.configure(image=self.attrPickerPI)

        self.setStatus()

    def setTile(self, tile):
        self.curTile = tile
        self.setAttribute(self.curAttribute)

    def setStatus(self):
        if self.status is None:
            return
        label = "tile $%02x attr %d" % (self.curTile, self.curAttribute)
        self.status.configure(text=label)

    def tilePickerCallback(self, event):
        if event.x >= 0 and event.x < 128 and event.y >= 0 and event.y < 128:
            tileX = event.x // 8
            tileY = event.y // 8
            newTileNo = tileY * 16 + tileX
            #print "mouse was clicked on tile", newTileNo
            self.setTile(newTileNo)
            return
        print "mouse was clicked at (%d, %d)" % (event.x, event.y)

    def attrPickerCallback(self, event):
        if event.x >= 0 and event.x < 128:
            attr = event.x // 32
            #print "mouse was clicked on attribute", attr
            self.setAttribute(attr)
            return
        print "mouse was clicked at (%d, %d)" % (event.x, event.y)

class NamDisplay(Canvas):
    def __init__(self, parent, doc, **kw):
        kw['width'] = 512
        kw['height'] = 480
        kw['relief']='raised'
        kw['highlightthickness'] = 0
        apply(Canvas.__init__, (self, parent), kw)
        self.doc = doc
        self.tile = []
        im = PILImage.new('RGB', (32, 32))
        for y in range(15):
            row = []
            for x in range(16):
                tile = ImageTk.PhotoImage(im)
                if True or ((x ^ y) & 1) == 0:
                    self.create_image(x * 32, y * 32, image=tile, anchor=NW)
                row.append(tile)
            self.tile.append(row)
        self.updating = False
        self.updScreen()

    def updScreen(self):
        self.updating = True
        for y in range(0, 30, 2):
            for x in range(0, 32, 2):
                self.updTile(x, y)
        self.updating = False
        self.update_idletasks()

    def updTile(self, x, y):
        if x < 0 or x >= 32 or y < 0 or y >= 30:
            return
        y = y & ~1
        x = x & ~1
        im = PILImage.new('RGB', (32, 32))
        dst = self.tile[y >> 1][x >> 1]
        for y1 in range(2):
            for x1 in range(2):
                (tileNo, attrNo) = self.doc.getTile(x + x1, y + y1)
                tile = self.doc.renderTile(tileNo, attrNo).resize((16, 16))
                tile.putpalette(self.doc.clut)
                im.paste(tile, (x1 * 16, y1 * 16))
        dst.paste(im)
        if not self.updating:
            self.update_idletasks()

class PackBits():
    def __init__(self, toWrite=''):
        self.bytes = array('b')
        self.closed = False
        self.mode = 'wb'
        self.name = '<PackBits>'
        self.newlines = None
        if toWrite:
            self.write(toWrite)

    def close(self):
        self.bytes = None
        self.closed = True

    def write(self, s):
        """Add a string to the buffer."""
        if not self.closed:
            self.bytes.fromstring(s)

    def tell(self):
        return len(self.bytes)

    def truncate(self, length):
        if not self.closed:
            del self[length:]

    def writelines(self, seq):
        """Add a sequence of strings to the buffer."""
        self.write(''.join(seq))

    def flush(self):
        """Compress the data to a file."""
        i = 0
        base = 0
        out = array('b')
        while base < len(self.bytes):

            # measure the run starting at t
            i = 1
            imax = min(128, len(self.bytes) - base)
            basebyte = self.bytes[base]
            while (i < imax
                   and basebyte == self.bytes[base + i]):
                i += 1
            # if the run is either length 3 or to the end of the file,
            # write it
            if i > 2 or base + i == len(self.bytes):
                out.append(1 - i)
                out.append(self.bytes[base])
                base += i
                continue

            # measure the nonrun starting at t
            i = 1
            imax = min(128, len(self.bytes) - base)
            while (i < imax
                   and (base + i + 2 >= len(self.bytes)
                        or self.bytes[base + i] != self.bytes[base + i + 1]
                        or self.bytes[base + i] != self.bytes[base + i + 2])):
                i += 1
            out.append(i - 1)
            out.extend(self.bytes[base:base + i])
            base += i
        return out

    @staticmethod
    def test():
        pb = PackBits('stopping stoppping stopppppi')
        data = pb.flush()
        print repr(data)

class UnPackBits(PackBits):
    def flush(self):
        out = array('b')
        base = 0
        while base < len(self.bytes):
            c = self.bytes[base]
            if c > 0 and c <= 127:
                b = self.bytes[base + 1]
                out.extend(self.bytes[base + 1:base + c + 2])
                base += 2 + c
            elif c >= -127:
                b = self.bytes[base + 1]
                out.fromlist([b] * (1 - c))
                base += 2
        return out

    @staticmethod
    def test():
        start = 'stopping stoppping stopppppi'
        packed = PackBits(start).flush().tostring()
        print repr(packed)
        unpacked = UnPackBits(packed).flush().tostring()
        print repr(unpacked)
        print "pass" if start == unpacked else "fail"

class App:
    filetypes = [
        ('NES nametable', '*.nam'),
        ('NES compressed nametable', '*.pkb'),
        ('PNG image', '*.png'),
        ('GIF image', '*.gif'),
        ('Windows bitmap', '*.bmp')
    ]

    def __init__(self, w, doc):
        import sys
        self.window = w
        self.doc = doc
        mbardata = [
            ("File", [
                ("&New Nametable", lambda: self.file_new_nam(), "Ctrl+N"),
                ("&Open Nametable...", lambda: self.file_open_nam(), "Ctrl+O"),
                ("Open &Pattern Table...", lambda: self.file_open_chr(), "Ctrl+L"),
                '-',
                ("&Save", lambda: self.file_save_nam(), "Ctrl+S"),
                ("Save &As...", lambda: self.file_save_nam_as(), "Ctrl+A"),
                '-',
                ("E&xit", lambda: self.file_quit(), "Ctrl+Q")
            ]),
            ("Help", [
                ("&About...", lambda: self.about())
            ])
        ]
        (menubar, menus) = build_menubar(w, mbardata)
        w.bind("<Control-n>", lambda e: self.file_new_nam())
        w.bind("<Control-N>", lambda e: self.file_new_nam())
        w.bind("<Control-o>", lambda e: self.file_open_nam())
        w.bind("<Control-O>", lambda e: self.file_open_nam())
        w.bind("<Control-l>", lambda e: self.file_open_chr())
        w.bind("<Control-L>", lambda e: self.file_open_chr())
        w.bind("<Control-s>", lambda e: self.file_save_nam())
        w.bind("<Control-S>", lambda e: self.file_save_nam())
        w.bind("<Control-q>", lambda e: self.file_quit())
        w.bind("<Control-Q>", lambda e: self.file_quit())

        self.tilePicker = TilePicker(w, doc)
        self.tilePicker.grid(row=0,column=0, sticky=NW)
        self.namDisplay = NamDisplay(w, doc, borderwidth=0)
        self.namDisplay.grid(row=0,column=1)
        self.namDisplay.bind("<Control-Button-1>", self.namPickupCallback)
        self.namDisplay.bind("<Control-B1-Motion>", self.namPickupCallback)
        self.namDisplay.bind("<Button-1>", self.namWriteCallback)
        self.namDisplay.bind("<B1-Motion>", self.namWriteCallback)
        w.wm_resizable(0,0)
        self.updWindowTitle()

    def namPickupCallback(self, event):
        if event.x >= 0 and event.x < 512 and event.y >= 0 and event.y < 512:
            x = event.x // 16
            y = event.y // 16
            (tile, attribute) = self.doc.getTile(x, y)
            self.tilePicker.curTile = tile
            self.tilePicker.setAttribute(attribute)
            return

    def namWriteCallback(self, event):
        if event.x >= 0 and event.x < 512 and event.y >= 0 and event.y < 512:
            x = event.x // 16
            y = event.y // 16
            t = self.tilePicker.curTile
            a = self.tilePicker.curAttribute
            self.doc.setTile(x, y, t, a)
            if not self.doc.unsaved:
                self.doc.setUnsaved(True)
                self.updWindowTitle()
            self.namDisplay.updTile(x, y)
            return

    def updWindowTitle(self):
        nfn = self.doc.namfilename
        if nfn is None:
            nfn = 'untitled'
        if self.doc.unsaved:
            nfn = '*' + nfn
        appname = '8name II'
        title = ' - '.join((nfn, appname))
        self.window.title(title)

    def file_new_nam(self):
        print "File > New Nametable"

    def file_open_nam(self):
        from tkFileDialog import askopenfilename
        filename = askopenfilename(parent=root,
                                   filetypes=self.filetypes,
                                   initialfile=self.doc.namfilename,
                                   title="Open Nametable")
        print "file open nam: filename is", filename
        if not isinstance(filename, basestring):
            return
        self.doc.loadnam(filename)
        self.namDisplay.updScreen()
        self.updWindowTitle()

    def file_open_chr(self):
        from tkFileDialog import askopenfilename
        filename = askopenfilename(parent=root,
                                   filetypes=[('Pattern Table', '*.chr')],
                                   initialfile=self.doc.namfilename,
                                   title="Open Pattern Table")
        if not isinstance(filename, str):
            return
        self.doc.loadchr(filename)
        self.tilePicker.updateWidgets()
        self.namDisplay.updScreen()

    def file_save_nam(self):
        if self.doc.namfilename is None:
            return self.file_save_nam_as()
        self.doc.savenam()
        self.updWindowTitle()

    def file_save_nam_as(self):
        from tkFileDialog import asksaveasfilename
        filename = asksaveasfilename(parent=root,
                                     filetypes=self.filetypes,
                                     title="Save Nametable As")
        ext = filename[-4:].lower()
        if ext in ('.png', '.gif', '.bmp'):
            print "Would save image to", filename
        else:
            self.doc.savenam(filename)
        self.updWindowTitle()

    def file_quit(self):
        self.window.destroy()

root = Tk()
app = App(root, NamFile('../spritecans.chr'))
root.mainloop()
print "remain:"
print "1. implement image saving"
print "2. implement and test loading"
print "3. implement and test compressed pkb support"
print "4. Implement stub methods for File menu items"
print "5. Warn on closing document where doc.unsaved is not False"
print "6. Write palette editor"


# Write a Python program to find out the number of CPUs using.
import multiprocessing
print(multiprocessing.cpu_count())


from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
import os

def main(train):
    if train==True:
        model = Unet(
            dim = 128,
            dim_mults = (1, 2, 2, 2)
        ).cuda()

        diffusion = GaussianDiffusion(
            model,
            image_size = 32,
            timesteps = 1000,   # number of steps
            loss_type = 'l2'    # L1 or L2
        ).cuda()

        trainer = Trainer(
            diffusion,
            '/home/congen/code/geoml_gan/data/cifar10',
            train=True,
            dataset_name='cifar10',
            image_size=32,
            train_batch_size = 64,
            train_lr = 2e-4,
            train_num_steps = 500001,         # total training steps
            gradient_accumulate_every = 2,    # gradient accumulation steps
            ema_decay = 0.9999,                # exponential moving average decay
            fp16 = True                   # turn on mixed precision training with apex
        )
        #trainer.load(20)
        trainer.train()
    else:

        model = Unet(
            dim=128,
            dim_mults=(1, 2, 2, 2)
        ).cuda()

        diffusion = GaussianDiffusion(
            model,
            image_size=32,
            timesteps=1000,  # number of steps
            loss_type='l2'  # L1 or L2
        ).cuda()

        trainer = Trainer(
            diffusion,
            '/home/congen/code/geoml_gan/data/cifar10',
            train=False,
            dataset_name='cifar10',
            image_size=32,
            train_batch_size=64,
            train_lr=2e-4,
            train_num_steps=200001,  # total training steps
            gradient_accumulate_every=2,  # gradient accumulation steps
            ema_decay=0.9999,  # exponential moving average decay
            fp16=True  # turn on mixed precision training with apex
        )
        trainer.test()

"""
    Usage:

        export CUDA_VISIBLE_DEVICES=2
        export PORT=6006
        export CUDA_HOME=/opt/cuda/cuda-10.2
        export TIME_STR=1
        python train.py

    :return:
    """
if __name__ == '__main__':
    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
        os.environ['CUDA_VISIBLE_DEVICES'] = '2'
    main(False)


# get human_ebv_tpms.py
import pandas as pd
import argparse
import os
import math
import datetime
import subprocess

# get basename from a file and path string
def get_basename(filepath):
    import os
    return os.path.basename(os.path.splitext(filepath)[0])

# get and format output directory
def format_odir(odir):
    import os
    cwd = os.getcwd()

    if odir != '':
        # if first character is not /, use cwd to make this an absolute path
        if odir[0] != '/' and odir[0] != '~':
            odir = cwd+odir
        if odir[-1] != '/':
            odir += '/'
    return odir

# make a dated output directory for the files used for the tracks
def make_dated_folder(odir, bname):
    date = datetime.datetime.now()
    date = date.strftime('%y%m%d')
    odir = odir+date+'_'+bname+'_figures/'

    if not os.path.isdir(odir):
        print('Making output directory '+odir)
        os.makedirs(odir)
    return odir

# get value associated with keyword in the 9th column of gtf
def get_field_value(key, fields):
    if key not in fields:
        return None
    else:
        return fields.split(key+' "')[1].split()[0].replace('";','')

# calculate tpm for a column in the abundance table
def get_tpm(df, col):
	new_col = 'TPM_'+col
	total_reads = df[d].sum()
	df[new_col] = df.apply(lambda x: float(x[d]*1000000)/total_reads, axis=1)

	return new_col, df

# calculate tpm for a column in the abundance table
def get_log_tpm(df, col, gene):
	tpm_col = 'TPM_'+col
	if not gene:
		new_col = 'log_'+tpm_col
	else:
		new_col = 'gene_log_'+TPM_col
	df[new_col] = df.apply(lambda x: math.log2(x[tpm_col]+1), axis=1)

	return new_col, df

# get gtf file name
parser = argparse.ArgumentParser(description='removes EBV transcripts from GTF file')
parser.add_argument('--human_gtf', help='GTF with human and EBV data')
parser.add_argument('--human_filt_ab', help='Filtered abundance file with human and EBV data')
parser.add_argument('--human_ab', help='Unfiltered abundance file with human and EBV data')
parser.add_argument('--ebv_filt_ab', help='EBV only filtered abundance file')
parser.add_argument('--ebv_ab', help='EBV only unfiltered abundance file')
parser.add_argument('--datasets', help='Comma-separated list of dataset names to use for human+ebv data')
parser.add_argument('--o', help='Prefix for output file')

args = parser.parse_args()

full_gtf = args.human_gtf
full_ab = args.human_filt_ab
full_unf_ab = args.human_ab
ebv_ab = args.ebv_filt_ab
ebv_unf_ab = args.ebv_ab

my_datasets = args.datasets.split(',')

oprefix = args.o

# get all human transcript ids
infile = open(full_gtf, 'r')
human_tids = []
ebv_tids = []
for i, line in enumerate(infile):
    line = line.replace('\n', '')
    temp = line.split('\t')
    fields = temp[-1]

    if temp[0] != 'chrEBV' and temp[2] == 'transcript':
        human_tids.append(get_field_value('talon_transcript', fields))
    elif temp[0] == 'chrEBV' and temp[2] == 'transcript':
        ebv_tids.append(get_field_value('talon_transcript', fields))

full_df = pd.read_csv(full_ab, sep='\t')
ebv_df = pd.read_csv(ebv_ab, sep='\t')

# reformat human table
# dc_datasets = ['D4', 'D5', 'D10', 'D11']
datasets = my_datasets
# full_df.drop(dc_datasets, inplace=True, axis=1) # drop datasets we don't want
full_df = full_df.loc[full_df[datasets].sum(axis=1) != 0] # drop transcripts with no reads in datasets we do want
full_df = full_df.loc[full_df['transcript_ID'].isin(human_tids)] # drop ebv transcripts
full_df['ebv'] = 'Human' # give human/ebv designation
full_df = full_df.loc[full_df.transcript_novelty != 'Genomic']

# drop genomic transcripts from the ebv dataset (b/c it's not pre-filtered)
ebv_df = ebv_df.loc[ebv_df.transcript_novelty != 'Genomic']
ebv_df['ebv'] = 'EBV' # human/ebv designation

# merge transcript df so TPMs can be calculated correctly
t_df = pd.concat([full_df, ebv_df])

# combine datasets
combine_datasets = True
if combine_datasets:
    t_df['combined'] = t_df[my_datasets].sum(axis=1)
    datasets = ['combined']

# # make sure the concatenation worked
# print(t_df.loc[t_df['transcript_ID'] == 121].head())
# print(ebv_df.loc[ebv_df['transcript_ID'] == 121].head())
# print(full_df.loc[full_df['transcript_ID'] == 121].head())

# get tpms and number of human transcripts for
# each dataset and for full/ebv
for d in datasets:

    # raw TPM
    tpm, t_df = get_tpm(t_df, d)

    # log2TPM
    log, t_df = get_log_tpm(t_df, d, 0)

    # sanity check - sum of all TPMs for each sample
    print('TPM total for {}: {}'.format(d, str(t_df[tpm].sum())))

    human_df = t_df.loc[(t_df[d] != 0) & (t_df['ebv'] == 'Human')]
    ebv_df = t_df.loc[(t_df[d] != 0) & (t_df['ebv'] == 'EBV')]
    n_human_transcripts = len(human_df.index)
    n_ebv_transcripts = len(ebv_df.index)
    print('Number of human transcripts in {}: {}'.format(d, str(n_human_transcripts)))
    print('Number of EBV transcripts in {}: {}'.format(d, str(n_ebv_transcripts)))

    # add columns for number of dataset human/ebv transcripts
    n_transcripts_col = 'n_'+d
    t_df[n_transcripts_col] = t_df.apply(lambda x:\
     n_human_transcripts if x['ebv'] == 'Human' else n_ebv_transcripts, axis=1)

    # add heights geom_text locations for dataset/human/ebv transcripts
    human_height = t_df.loc[t_df.ebv == 'Human'][log].max()+1
    ebv_height = t_df.loc[t_df.ebv == 'EBV'][log].max()+1
    height_col = d+'_height'
    t_df[height_col] = t_df.apply(lambda x:\
     human_height if x.ebv == 'Human' else ebv_height, axis=1)

    # print(human_height)
    # print(ebv_height)

# print(t_df.head())
# print(t_df.tail())

# write gene and transcript tables to a csv
t_df['dot_size'] = t_df.apply(lambda x: 1 if x['ebv'] == 'EBV' else 0.6, axis=1)
t_df['alpha'] = t_df.apply(lambda x: 0.5 if x['ebv'] == 'EBV' else 0.2, axis=1)

# bname = get_basename(ebv_ab)
# odir = format_odir(os.path.dirname(ebv_ab))
# odir = make_dated_folder(odir,bname)
to = oprefix+'_ebv_human_transcript_abundance.csv'
t_df.to_csv(to, sep=',', index=False)

## get transcript tpms without filtering for bioreps to use for gene tpms
# read in the unfiltered datasets
full_df = pd.read_csv(full_unf_ab, sep='\t')
ebv_df = pd.read_csv(ebv_unf_ab, sep='\t')

# reformat human table
# dc_datasets = ['D4', 'D5', 'D10', 'D11']
datasets = my_datasets
# full_df.drop(dc_datasets, inplace=True, axis=1) # drop datasets we don't want
full_df = full_df.loc[full_df['transcript_ID'].isin(human_tids)] # drop ebv transcripts
full_df['ebv'] = 'Human' # give human/ebv designation
full_df = full_df.loc[full_df.transcript_novelty != 'Genomic']

# drop genomic transcripts from the ebv dataset (b/c it's not pre-filtered)
ebv_df = ebv_df.loc[ebv_df.transcript_novelty != 'Genomic']
ebv_df['ebv'] = 'EBV' # human/ebv designation

# merge transcript df so TPMs can be calculated correctly
t_df = pd.concat([full_df, ebv_df])

# combine datasets
combine_datasets = True
if combine_datasets:
    t_df['combined'] = t_df[datasets].sum(axis=1)
    datasets = ['combined']

# # make sure the concatenation worked
# print(t_df.loc[t_df['transcript_ID'] == 121].head())
# print(ebv_df.loc[ebv_df['transcript_ID'] == 121].head())
# print(full_df.loc[full_df['transcript_ID'] == 121].head())

# get tpms and number of human transcripts for
# each dataset and for full/ebv
for d in datasets:

    # raw TPM
    tpm, t_df = get_tpm(t_df, d)

    # log2TPM
    log, t_df = get_log_tpm(t_df, d, 0)

    # sanity check - sum of all TPMs for each sample
    print('TPM total for {}: {}'.format(d, str(t_df[tpm].sum())))

    human_df = t_df.loc[(t_df[d] != 0) & (t_df['ebv'] == 'Human')]
    ebv_df = t_df.loc[(t_df[d] != 0) & (t_df['ebv'] == 'EBV')]
    n_human_transcripts = len(human_df.index)
    n_ebv_transcripts = len(ebv_df.index)
    print('Number of human transcripts in {}: {}'.format(d, str(n_human_transcripts)))
    print('Number of EBV transcripts in {}: {}'.format(d, str(n_ebv_transcripts)))

    # add columns for number of dataset human/ebv transcripts
    n_transcripts_col = 'n_'+d
    t_df[n_transcripts_col] = t_df.apply(lambda x:\
     n_human_transcripts if x['ebv'] == 'Human' else n_ebv_transcripts, axis=1)

    # add heights geom_text locations for dataset/human/ebv transcripts
    human_height = t_df.loc[t_df.ebv == 'Human'][log].max()+1
    ebv_height = t_df.loc[t_df.ebv == 'EBV'][log].max()+1
    height_col = d+'_height'
    t_df[height_col] = t_df.apply(lambda x:\
     human_height if x.ebv == 'Human' else ebv_height, axis=1)

# get gene tpms
cols = []
for d in datasets:
    cols.append(d)
    cols.append('TPM_'+d)
g_df = t_df.groupby(['gene_ID', 'gene_novelty', 'ebv'])[cols].sum()
g_df.reset_index(inplace=True)

# # make sure the groupby worked
# print(g_df.loc[g_df['gene_ID'] == 16].head())
# print(t_df.loc[t_df['gene_ID'] == 16].head())
# print(t_df.loc[t_df['gene_ID'] == 16].head())

# get tpms, heights, and numbers for gene
for d in datasets:
    # log2TPM
    log, g_df = get_log_tpm(g_df, d, 0)

    human_df = g_df.loc[(g_df[d] != 0) & (g_df['ebv'] == 'Human')]
    ebv_df = g_df.loc[(g_df[d] != 0) & (g_df['ebv'] == 'EBV')]
    n_human_genes = len(human_df.index)
    n_ebv_genes = len(ebv_df.index)
    print('Number of human genes in {}: {}'.format(d, str(n_human_genes)))
    print('Number of EBV genes in {}: {}'.format(d, str(n_ebv_genes)))

    # add columns for number of dataset human/ebv genes
    n_genes_col = 'n_'+d
    g_df[n_genes_col] = g_df.apply(lambda x:\
     n_human_genes if x['ebv'] == 'Human' else n_ebv_genes, axis=1)

    # add heights geom_text locations for dataset/human/ebv transcripts
    human_height = g_df.loc[g_df.ebv == 'Human'][log].max()+0.4
    ebv_height = g_df.loc[g_df.ebv == 'EBV'][log].max()+0.4
    height_col = d+'_height'
    g_df[height_col] = g_df.apply(lambda x:\
     human_height if x.ebv == 'Human' else ebv_height, axis=1)

    print(human_height)
    print(ebv_height)

print(g_df.head())
print(g_df.tail())

# add different dot sizes for human/ebv
# t_df['dot_size'] = t_df.apply(lambda x: 1 if x['ebv'] == 'EBV' else 0.6, axis=1)
g_df['dot_size'] = g_df.apply(lambda x: 1 if x['ebv'] == 'EBV' else 0.6, axis=1)

# t_df['alpha'] = t_df.apply(lambda x: 0.5 if x['ebv'] == 'EBV' else 0.2, axis=1)
g_df['alpha'] = g_df.apply(lambda x: 0.5 if x['ebv'] == 'EBV' else 0.2, axis=1)

# # rename gene/transcript novelty columns
# t_df.rename(index=str, columns={\
#     'transcript_novelty':'Isoform Type'}, inplace=True)
# g_df.rename(index=str, columns={\
#     'gene_novelty':'Gene Type'}, inplace=True)

# write gene table to a csv

go = oprefix+'_ebv_human_gene_abundance.csv'
g_df.to_csv(go, sep=',', index=False)

# make graphs
cmd = 'Rscript plot_ebv_v_human_abundances.R --gene_csv {} --transcript_csv {}'\
	  ' --datasets {}'.format(go, to, ','.join(datasets))
# print(cmd)


#!/usr/bin/env python

"""

    Normalize site observed ancestry by genome-wide average.
    2*ID_average - ExpectedCopiesOfPop1Ancestry.
    @Author: wavefancy@gmail.com

    Usage:
        HapmixEPop1NormalizedByIDAverage.py -a IDaverage
        HapmixEPop1NormalizedByIDAverage.py -h | --help | -v | --version | -f | --format

    Notes:
        1. Read ExpectedCopiesOfPop1Ancestry from stdin, and output to stdout.

    Options:
        -a IDaverage     Individual average for pop1, one line one person.
        -h --help        Show this screen.
        -v --version     Show version.
        -f --format      Show input/output file format example.

"""
import sys
from docopt import docopt
from signal import signal, SIGPIPE, SIG_DFL

def ShowFormat():
    '''File format example'''
    print('''
    #ExpectedCopiesOfPop1Ancestry, one column each person.
    ------------------------
1 1 2 1 2 2
1 2 2 1 2 2

    #id average:
    ------------------------
0.8
0.5
0.1

    #output:
    ------------------------
0.40    0.00    -0.20
-0.60   0.00    -0.20
          ''');

if __name__ == '__main__':
    args = docopt(__doc__, version='1.0')
    # print(args)
    # sys.exit(0)

    if(args['--format']):
        ShowFormat()
        sys.exit(-1)

    idav = [] # 2*ID_average
    with open(args['-a'],'r') as ifile:
        for line in ifile:
            line = line.strip()
            if line:
                idav.append(2 * float(line))

    checkLen = True
    for line in sys.stdin:
        line = line.strip()
        if line:
            ss = line.split()
            ss = [float(x) for x in ss] #number of pop1 copy for each person.

            if checkLen:
                if len(ss) != len(idav):
                    sys.stderr.write('Error: numbr of individuals in ID_average file is not the same as that in sys.stdin.\n')
                    sys.exit(-1)
                else:
                    checkLen = False

            out = [ '%.2f'%(y-x) for x,y in zip(idav, ss)]
            sys.stdout.write('%s\n'%('\t'.join(out)))

sys.stdout.flush()
sys.stdout.close()
sys.stderr.flush()
sys.stderr.close()


# -*- coding: utf-8 -*-
# Copyright 2021 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for execution_util.py."""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

import subprocess
from unittest import mock

from gslib import exception
from gslib.tests import testcase
from gslib.utils import execution_util

class TestExecutionUtil(testcase.GsUtilUnitTestCase):
  """Test execution utils."""

  @mock.patch.object(subprocess, 'Popen')
  def testExternalCommandReturnsNoOutput(self, mock_Popen):
    mock_command_process = mock.Mock()
    mock_command_process.returncode = 0
    mock_command_process.communicate.return_value = (None, None)
    mock_Popen.return_value = mock_command_process

    stdout, stderr = execution_util.ExecuteExternalCommand(['fake-command'])
    self.assertIsNone(stdout)
    self.assertIsNone(stderr)

    mock_Popen.assert_called_once_with(['fake-command'],
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)

  @mock.patch.object(subprocess, 'Popen')
  def testExternalCommandReturnsStringOutput(self, mock_Popen):
    mock_command_process = mock.Mock()
    mock_command_process.returncode = 0
    mock_command_process.communicate.return_value = ('a', 'b')
    mock_Popen.return_value = mock_command_process

    stdout, stderr = execution_util.ExecuteExternalCommand(['fake-command'])
    self.assertEqual(stdout, 'a')
    self.assertEqual(stderr, 'b')

    mock_Popen.assert_called_once_with(['fake-command'],
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)

  @mock.patch.object(subprocess, 'Popen')
  def testExternalCommandReturnsBytesOutput(self, mock_Popen):
    mock_command_process = mock.Mock()
    mock_command_process.returncode = 0
    mock_command_process.communicate.return_value = (b'a', b'b')
    mock_Popen.return_value = mock_command_process

    stdout, stderr = execution_util.ExecuteExternalCommand(['fake-command'])
    self.assertEqual(stdout, 'a')
    self.assertEqual(stderr, 'b')

    mock_Popen.assert_called_once_with(['fake-command'],
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)

  @mock.patch.object(subprocess, 'Popen')
  def testExternalCommandReturnsNoOutput(self, mock_Popen):
    mock_command_process = mock.Mock()
    mock_command_process.returncode = 1
    mock_command_process.communicate.return_value = (None, b'error')
    mock_Popen.return_value = mock_command_process

    with self.assertRaises(exception.ExternalBinaryError):
      execution_util.ExecuteExternalCommand(['fake-command'])

    mock_Popen.assert_called_once_with(['fake-command'],
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)

  @mock.patch.object(subprocess, 'Popen')
  def testExternalCommandRaisesFormattedStderr(self, mock_Popen):
    mock_command_process = mock.Mock()
    mock_command_process.returncode = 1
    mock_command_process.communicate.return_value = (None, b'error.\n')
    mock_Popen.return_value = mock_command_process

    with self.assertRaisesRegexp(exception.ExternalBinaryError, 'error'):
      execution_util.ExecuteExternalCommand(['fake-command'])