#!/usr/bin/env python

from pygments.lexers.c_cpp import CppLexer
from pygments.token import Keyword, Name, Operator


class CustomLexer(CppLexer):
    name = "Pro"
    aliases = ["pro"]
    filenames = ["*.pro"]
    EXTRA_KEYWORDS = [
        "Macro",
        "Group",
        "Function",
        "Constraint",
        "FunctionSpace",
        "Jacobian",
        "Integration",
        "Formulation",
        "Resolution",
        "PostProcessing",
        "PostOperation",
    ]

    OTHER = [
        "Case",
        "Equation",
        "BasisFunction",
        "Quantity",
        "SubSpace",
        "GlobalQuantity",
        "Operation",
        "System",
        "UsingPost",
    ]
    CONSTS = ["Pi", "0D", "1D", "2D", "3D", "All"]

    CTRL = ["If", "ElseIf", "Else", "EndIf", "Return"]

    CALLS = ["Call", "Function", "Macro"]

    GROUPS = [
        "Region",
        "Global",
        "NodesOf",
        "EdgesOf",
        "FacetsOf",
        "VolumesOf",
        "ElementsOf",
        "GroupsOfNodesOf",
        "GroupsOfEdgesOf",
        "GroupsOfEdgesOnNodesOf",
        "EdgesOfTreeIn",
        "FacetsOfTreeIn",
        "DualNodesOf",
        "DualEdgesOf",
        "DualFacetsOf",
        "DualVolumesOf",
    ]

    KW = [
        "Analytic",
        "Branch",
        "DefineConstant",
        "DefineFunction",
        "DefineGroup",
        "DestinationSystem",
        "Entity",
        "EntitySubType",
        "EntityType",
        "GeoElement",
        "In",
        "Include",
        "IndexOfSystem",
        "InterpolationBilinear",
        "List",
        "ListAlt",
        "ListFromFile",
        "Loop",
        "Name",
        "NameOfBasisFunction",
        "NameOfFormulation",
        "NameOfPostProcessing",
        "NameOfCoef",
        "NameOfConstraint",
        "NameOfMesh",
        "NameOfResolution",
        "NameOfSpace",
        "Not",
        "NumberOfPoints",
        "OnOneSideOf",
        "OriginSystem",
        "Solver",
        "StartingOn",
        "StrCat",
        "SubRegion",
        "Support",
        "Symmetry",
        "TimeFunction",
        "RegionRef",
        "Coefficient",
        "Type",
        "Dof",
        "BF",
        "d",
        "Grad",
        "Curl",
        "Rot",
        "Div",
        "dInv",
        "GradInv",
        "CurlInv",
        "RotInv",
        "DivInv",
        "ComplexValue",
    ]

    FUNC = [
        "Hypot",
        "Rand",
        "Fmod",
        "Fabs",
        "Floor",
        "Exp",
        "Ceil",
        "Modulo",
        "Log",
        "Log10",
        "ArcCos",
        "Acos",
        "Asin",
        "ArcSin",
        "Atan",
        "ArcTan",
        "Atan2",
        "ArcTan2",
        "Ceil",
        "Cos",
        "Cosh",
        "Floor",
        "Hypot",
        "Rand",
        "Round",
        "Sqrt",
        "Sin",
        "Sinh",
        "Tan",
        "Tanh",
        "Cross",
        "Hypot",
        "Norm",
        "SquNorm",
        "Unit",
        "Transpose",
        "TTrace",
        "F_Cos_wt_p",
        "F_Sin_wt_p",
        "F_Period",
        "Laplace",
        "GradLaplace",
        "Helmholtz",
        "GradHelmholtz",
        "Complex",
        "Re",
        "Im",
        "Vector",
        "Tensor",
        "TensorV",
        "TensorSym",
        "TensorDiag",
        "CompX",
        "CompY",
        "CompZ",
        "CompXX",
        "CompXY",
        "CompXZ",
        "CompYX",
        "CompYY",
        "CompYZ",
        "CompZX",
        "CompZY",
        "CompZZ",
        "X",
        "Y",
        "Z",
        "XYZ",
    ]

    CMD = [
        "Printf",
        "Normal",
        "NormalSource",
        "F_CompElementNum",
        "InterpolationLinear",
        "dInterpolationLinear",
        "InterpolationAkima",
        "dInterpolationAkima",
        "Order",
        "Assign",
        "Init",
        "AssignFromResolution",
        "InitFromResolution",
        "Network",
        "Link",
        "LinkCplx",
        "Form0",
        "Form1",
        "Form2",
        "Form3",
        "Form1P",
        "Form2P",
        "Scalar",
        "Vector",
        "BF_Node_2E",
        "BF_Node",
        "BF_Edge",
        "BF_Facet",
        "BF_Volume",
        "BF_GradNode",
        "BF_CurlEdge",
        "BF_DivFacet",
        "BF_GroupOfNodes",
        "BF_GradGroupOfNodes",
        "BF_GroupOfEdges",
        "BF_CurlGroupOfEdges",
        "BF_PerpendicularEdge",
        "BF_CurlPerpendicularEdge",
        "BF_GroupOfPerpendicularEdge",
        "BF_CurlGroupOfPerpendicularEdge",
        "BF_PerpendicularFacet",
        "BF_DivPerpendicularFacet",
        "BF_Region",
        "BF_RegionX",
        "BF_RegionY",
        "BF_RegionZ",
        "BF_Global",
        "BF_dGlobal",
        "BF_NodeX",
        "BF_NodeY",
        "BF_NodeZ",
        "BF_Zero",
        "BF_One",
        "AliasOf",
        "AssociatedWith",
        "Vol",
        "Sur",
        "Lin",
        "VolAxi",
        "SurAxi",
        "VolAxiSqu",
        "VolSphShell",
        "VolAxiSphShell",
        "VolAxiSquSphShell",
        "VolRectShell",
        "VolAxiRectShell",
        "VolAxiSquRectShell",
        "Gauss",
        "GaussLegendre",
        "Line",
        "Triangle",
        "Quadrangle",
        "Tetrahedron",
        "Hexahedron",
        "Prism",
        "Pyramid",
        "Point",
        "FemEquation",
        "Galerkin",
        "deRham",
        "Local",
        "Global",
        "Integral",
        "Dt",
        "DtDof",
        "DtDt",
        "DtDtDof",
        "JacNL",
        "NeverDt",
        "Evaluate",
        "Generate",
        "Solve",
        "SolveAgain",
        "EigenSolve",
        "GenerateRHSGroup",
        "EigenvalueLegend",
        "GenerateJac",
        "SolveJac",
        "GenerateSeparate",
        "Update",
        "InitSolution",
        "SaveSolution",
        "SaveSolutions",
        "TransferSolution",
        "TransferInitSolution",
        "SetTime",
        "SetFrequency",
        "SystemCommand",
        "If",
        "EndIf",
        "Else",
        "Print",
        "Lanczos",
        "FourierTransform",
        "TimeLoopTheta",
        "TimeLoopNewmark",
        "IterativeLoop",
        "OnElementsOf",
        "OnRegion",
        "OnGlobal",
        "OnSection",
        "OnGrid",
        "OnPoint",
        "OnLine",
        "OnPlane",
        "OnBox",
        "File",
        "Depth",
        "Skin",
        "Smoothing",
        "HarmonicToTime",
        "Dimension",
        "TimeStep",
        "Frequency",
        "Format",
        "Adapt",
        "Target",
        "Value",
        "Sort",
        "Iso",
        "NoNewLine",
        "ChangeOfCoordinates",
        "ChangeOfValues",
        "Gmsh",
        "GmshOpen",
        "GmshParsed",
        "SimpleTable",
        "NodeTable",
        "Table",
        "TimeTable",
        "Gnuplot",
        "Adaptation",
    ]

    def get_tokens_unprocessed(self, text):
        for index, token, value in CppLexer.get_tokens_unprocessed(self, text):
            # print(token)
            if token is Name and value in self.EXTRA_KEYWORDS:
                yield index, Name.Class, value
            elif token is Name and value in self.OTHER:
                yield index, Keyword.Type, value
            elif token is Name and value in self.CONSTS:
                yield index, Name.Builtin, value
            elif token is Name and value in self.CTRL:
                yield index, Name.Tag, value
            elif token is Name and value in self.CALLS:
                yield index, Operator.Word, value
            elif token is Name and value in self.GROUPS:
                yield index, Name.Attribute, value
            elif token is Name and value in self.KW:
                yield index, Keyword.Pseudo, value
            elif token is Name and value in self.FUNC:
                yield index, Name.Function, value
            elif token is Name and value in self.CMD:
                yield index, Operator.Word, value
            else:
                yield index, token, value
