import os
import tempfile
import pytest
import warnings
import numpy as np
import onnxruntime as ort
import torch
from torch import nn as nn
from typing import Optional, Union, Tuple, List
import transformer_engine.pytorch as te
from transformer_engine.common import recipe
import transformer_engine_extensions as tex
from transformer_engine.pytorch.cpp_extensions import gemm, fp8_gemm, gelu, cast_to_fp8, cast_from_fp8
from transformer_engine.pytorch.module.base import get_workspace
import transformer_engine.pytorch.cpp_extensions as texcpp
import transformer_engine.pytorch.softmax as softmax_defs
from transformer_engine.pytorch.utils import get_default_init_method
from transformer_engine.pytorch.export import is_in_onnx_export_mode
from transformer_engine.pytorch.fp8 import FP8GlobalStateManager

# Global test configuration knobs.

# Enable this to serialize test inputs and outputs to file (as a Polygraphy RunResults instance).
SAVE_TEST_IO = bool(int(os.getenv("NVTE_ONNX_EXPORT_SAVE_TEST_IO", "0")))

if SAVE_TEST_IO:
    from polygraphy.json import save_json
    from polygraphy.comparator import RunResults

# The directory where generated ONNX test models are stored.
NVTE_TEST_ARTIFACTS_DIR = os.environ.get('NVTE_TEST_ARTIFACTS_DIR')
NVTE_TEST_ARTIFACTS_DIR = NVTE_TEST_ARTIFACTS_DIR or os.path.join(tempfile.gettempdir(), "./gen_onnx_models")

# The directory where this file is stored.
TESTS_DIR = os.path.dirname(os.path.abspath(__file__))

# ScaledUpperTriangMaskedSoftmax is exported via ONNX::Trilu which was introduced in opset 14.
TRILU_OPSET = 14
# Opset used in the ONNX files generated by the tests.
OPSET = 17
assert OPSET >= TRILU_OPSET

# Shared library implementing custom FP8 Q/DQ operators for ONNX Runtime (ORT).
ORT_CUSTOM_OPS_LIB = os.path.join(TESTS_DIR, "./libcustom_ort_fp8_qdq_ops.so")

fp8_available, reason_for_no_fp8 = FP8GlobalStateManager.is_fp8_available()
skip_FP8 = pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)

supported_activations = ["gelu", "relu", "reglu", "geglu", "swiglu"]

all_normalizations = ["LayerNorm", "RMSNorm"]

@pytest.fixture()
def seed_default_rng():
    """Reseed the PRNG for test reproducibility"""
    torch.random.seed()

@pytest.fixture()
def set_max_seq_len(max_seq_len=128):
    """Set the maximum sequence length that can be used for attention masking"""
    os.environ["NVTE_ONNX_KVCACHE_MAX_SEQ_LEN"] = f"{max_seq_len}"

def create_fp8_recipe():
    return recipe.DelayedScaling(margin=0, interval=1, fp8_format=recipe.Format.E4M3)

def do_export(
    model: torch.nn.Module,
    inp: torch.Tensor,
    fname: str,
    use_fp8: bool=True,
    opset: int=OPSET,
    input_names: List[str]=None,
    output_names: List[str]=None,
    dynamic_axes: List[str]=None
):
    """Export to ONNX"""
    fp8_recipe = create_fp8_recipe()
    input_names = input_names or ["input"]
    output_names = output_names or ["output"]

    with torch.inference_mode(), te.fp8_autocast(enabled=use_fp8, fp8_recipe=fp8_recipe), warnings.catch_warnings():
        warnings.filterwarnings(
            action='ignore',
            category=torch.jit.TracerWarning,
            module=r'.*'
        )

        model.cuda().eval()
        os.makedirs(NVTE_TEST_ARTIFACTS_DIR, exist_ok=True)
        fname = os.path.join(NVTE_TEST_ARTIFACTS_DIR, fname)

        inps = inp if isinstance(inp, list) or isinstance(inp, tuple) else (inp,)
        assert len(inps) == len(input_names)
        inds_to_del = [i for i in range(len(inps)) if inps[i] is None]
        input_names = [input_names[i] for i in range(len(inps)) if i not in inds_to_del]

        with te.onnx_export(True):
            torch.onnx.export(
                model,
                inps,
                fname,
                verbose=True,
                dynamic_axes=dynamic_axes,
                opset_version=opset,
                input_names=input_names,
                output_names=output_names,
                do_constant_folding=True,
                operator_export_type=torch.onnx.OperatorExportTypes.ONNX_FALLTHROUGH)

def to_numpy(tensor):
    if isinstance(tensor, torch.Tensor):
        if tensor.dtype == torch.bfloat16:
            tensor = tensor.type(torch.float32)
        tensor = tensor.detach().cpu().numpy()
    return tensor

def set_layer_scale(module: torch.nn.Module, scale: float, num_gemms: int):
    """Initialize the FP8 quantization scales in module"""
    NB_SCALES_PER_GEMM = 3  # One scale per: input, weights, and output GEMM tensors.
    nb_total_scales = num_gemms * NB_SCALES_PER_GEMM
    module.fp8_init(num_gemms)
    module.fp8_meta["scaling_fwd"].scale = torch.ones(
        nb_total_scales, dtype=torch.float32, device="cuda") / scale
    module.fp8_meta["scaling_fwd"].scale_inv = torch.ones(
        nb_total_scales, dtype=torch.float32, device="cuda") * scale

def te_infer(model: torch.nn.Module, inps: Union[Tuple[torch.tensor], torch.tensor], is_fp8: bool):
    """Transformer Engine forward propagation."""
    fp8_recipe = create_fp8_recipe()
    with torch.inference_mode(), te.fp8_autocast(enabled=is_fp8, fp8_recipe=fp8_recipe), warnings.catch_warnings():
        te_outputs = model(*inps if isinstance(inps, tuple) else (inps,))
        if not isinstance(te_outputs, tuple):
            te_outputs = (te_outputs,)
        return te_outputs

def compare_outputs(onnx_outputs, te_outputs, atol, rtol, max_errors_printed, allow_cnt_errors, fname):
    """ Compare ORT and TE outputs."""
    assert len(onnx_outputs) == len(te_outputs)
    # Compare ORT and PyTorch outputs.
    for onnx_output, te_output in zip(onnx_outputs, te_outputs):
        # np.isclose: abs(a - b) <= (atol + rtol * abs(b))
        te_output = to_numpy(te_output)
        onnx_output = to_numpy(onnx_output)
        ac = ~np.isclose(onnx_output, te_output, atol=atol, rtol=rtol)
        mismatches = ac.nonzero()
        mismatched_ids = [loc for loc in zip(*mismatches)]
        if mismatched_ids:
            # Log some information in case of error.
            print("*" * 100)
            nb_errors = len(mismatched_ids)
            nb_vals = min(nb_errors, max_errors_printed)
            print(f"Detected {nb_errors} diverging values (output shape={onnx_output.shape})")
            print(f"Showing first {nb_vals} errors (ONNX -- TE):")
            abs_err = np.abs(onnx_output - te_output)
            errors = abs_err[mismatches]
            for loc in mismatched_ids[:nb_vals]:
                ref = te_output[loc]
                print(f"{onnx_output[loc]} -- {te_output[loc]} err={abs_err[loc]} > {atol + rtol * abs(ref)}")
            print(f"Max error: {np.max(errors)}")
            if nb_errors > allow_cnt_errors:
                raise ValueError(f"Output validation of {fname} failed with {nb_errors} errors")

def serialize_inputs_outputs(
    fname: str,
    inputs: Union[Tuple[torch.Tensor], torch.Tensor],
    te_outputs: List[torch.Tensor],
    input_names: Optional[List[str]] = None,
    output_names: Optional[List[str]] = None,
):
    if not SAVE_TEST_IO:
        return

    fname = os.path.join(NVTE_TEST_ARTIFACTS_DIR, fname)

    input_names = input_names or ["input"]
    output_names = output_names or ["output"]
    inputs = inputs if isinstance(inputs, list) or isinstance(inputs, tuple) else (inputs,)
    named_inputs = zip(input_names, inputs)
    input_data = [{k: v.cpu() for k, v in named_inputs if v is not None}]
    json_fname = fname[:-len(".onnx")] + "_inputs.json"
    save_json(input_data, json_fname, description="custom input data")

    json_fname = fname[:-len(".onnx")] + "_output.json"
    named_outputs = zip(output_names, te_outputs)
    output_data = {k: v.detach().cpu() for k, v in named_outputs if v is not None}
    custom_outputs = RunResults()
    custom_outputs.add([output_data], runner_name="custom_runner")
    custom_outputs.save(json_fname)

def validate_result(
    fname: str,
    inps: Union[Tuple[torch.Tensor], torch.Tensor],
    model: torch.nn.Module,
    atol: float=1.e-8, # np.isclose default atol
    rtol: float=1.e-5, # np.isclose default rtol
    max_errors_printed: int=10,
    is_fp8: bool=False,
    allow_cnt_errors: int=0,
    input_names: List[str]=None,
    output_names: List[str]=None,
    te_outputs: List[torch.Tensor]=None,
):
    """Compare the outputs of a Transformer Engine (TE) module vs the outputs of its ONNX
    representation using ONNX Runtime (ORT) and ensure they are close.

    The purpose of the output comparison is to validate that TE models are converted to
    their correct ONNX representation by testing that TE and ORT outputs match within some
    small threshold (allowing for finite precision errors).

    Argument `allow_cnt_errors` reduces test failure noise due to spurious errors by ignoring,
    a very small number (0-3) of outliers. This is fine to do because these outliers are due to
    small kernel implementation differences between TE and ORT and do not imply an incorrect ONNX
    representation (the tests assume both ORT or TE kernels are correct).

    Argument `te_outputs` can be used to provide pre-computed TE outputs.
    """

    def create_ort_session(fname: str, is_fp8: bool):
        def load_custom_ops(session_opts: ort.SessionOptions):
            """For FP8 validation with ORT we need to load our custom FP8 Q/DQ extension."""
            if not os.path.exists(ORT_CUSTOM_OPS_LIB):
                raise FileNotFoundError(f"Unable to find {ORT_CUSTOM_OPS_LIB}")
            session_opts.register_custom_ops_library(ORT_CUSTOM_OPS_LIB)
            print("registered custom FP8 Q/DQ ops!")

        """Create an ONNX Runtime session for validation."""
        kwargs = {"providers": ['CUDAExecutionProvider', 'CPUExecutionProvider']}
        if is_fp8:
            sess_options = ort.SessionOptions()
            load_custom_ops(sess_options)
            kwargs["sess_options"] = sess_options

        s = ort.InferenceSession(fname, **kwargs)
        return s

    def create_ort_input_dict(session, inputs):
        inputs = inputs if isinstance(inputs, list) or isinstance(inputs, tuple) else (inputs,)
        input_names = [x.name for x in session.get_inputs()]
        inps = [to_numpy(x) for x in inputs if x is not None]
        inp_dict = dict(zip(input_names, inps))
        return inp_dict

    input_names = input_names or ["input"]
    output_names = output_names or ["output"]

    # Run ORT session and TE model.
    fname = os.path.join(NVTE_TEST_ARTIFACTS_DIR, fname)
    if not te_outputs:
        te_outputs = te_infer(model, inps, is_fp8)
    ort_s = create_ort_session(fname, is_fp8)
    input_feed = create_ort_input_dict(ort_s, inps)
    onnx_outputs = ort_s.run(None, input_feed=input_feed)
    compare_outputs(onnx_outputs, te_outputs, atol, rtol, max_errors_printed, allow_cnt_errors, fname)

def create_meta(scale_factor: float, size: int=1):
    meta = tex.FP8TensorMeta()
    meta.amax_history = torch.zeros(1, size, dtype=torch.float32, device="cuda")
    meta.scale_inv = torch.ones(size, dtype=torch.float32, device="cuda") / scale_factor
    meta.scale = torch.ones(size, dtype=torch.float32, device="cuda") * scale_factor
    return meta

def dtype2str(dtype: torch.dtype, fake_bf16_io=False):
    if fake_bf16_io:
        assert dtype == torch.bfloat16
        return "_fake_bf16"
    return {
        torch.float32: "_fp32",
        torch.float16: "_fp16",
        torch.bfloat16: "_bf16",
    }[dtype]

def as_te_type(dtype: torch.dtype):
    return {
        torch.float32: tex.DType.kFloat32,
        torch.float16: tex.DType.kFloat16,
        torch.bfloat16: tex.DType.kBFloat16,
    }[dtype]

def get_attn_mask_str(use_mask, attn_mask_type):
    # See FusedScaleMaskSoftmax::forward_fused_softmax for logic behind names.
    if attn_mask_type is None:
        return "_mask" if use_mask else "_no-mask"
    attn_mask_str = "_arbitrary-no-mask"
    attn_mask_str = "_causal-mask" if attn_mask_type == "causal" else attn_mask_str
    attn_mask_str = "_arbitrary-mask" if use_mask and attn_mask_type == "arbitrary" else attn_mask_str
    return attn_mask_str

"""
Tests cases begin here.
"""

@skip_FP8
@pytest.mark.parametrize("scale_factor", [1, 224])
@pytest.mark.parametrize(
    "precision,             atol", [
    [torch.float32,         1e-7],
    [torch.float16,         1e-7],
    [torch.bfloat16,        5e-3],
    ["fake-torch.bfloat16", 5e-3],
])
def test_export_cast_ops(seed_default_rng, scale_factor: float, atol: float, precision: torch.dtype):
    fake_bf16_io = precision == "fake-torch.bfloat16"
    # reset precision to torch.bfloat16 after capturing fake BF16 mode
    precision = torch.bfloat16 if precision == "fake-torch.bfloat16" else precision

    class TestFP8_QDQ(nn.Module):
        def __init__(self, fake_bf16_io):
            super().__init__()
            self.fp8_tensor = 0
            self.meta = create_meta(scale_factor)
            self.highprec_type = as_te_type(precision)
            self.fp8_type = tex.DType.kFloat8E4M3
            self.fake_bf16_io = fake_bf16_io

        def forward(self, inp):
            ret = cast_to_fp8(
                inp,
                self.meta,
                self.fp8_tensor,
                self.fp8_type)

            ret = cast_from_fp8(
                ret,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                self.highprec_type)
            if self.fake_bf16_io:
                ret = ret.type(torch.float32)
            return ret

    # Set dimensions (these are arbitrary).
    in_features = 64
    hidden_size = 256
    inp = torch.randn(hidden_size, in_features, device="cuda",
        dtype=torch.float if fake_bf16_io else precision)
    high_prec_str = dtype2str(precision, fake_bf16_io=fake_bf16_io)
    fname = f"te.cast_fp8_{scale_factor}{high_prec_str}.onnx"
    model = TestFP8_QDQ(fake_bf16_io)

    do_export(model, inp, fname)
    te_outputs = te_infer(model, inp, is_fp8=True)
    serialize_inputs_outputs(fname, inp, te_outputs)
    if fake_bf16_io or precision != torch.bfloat16:
        validate_result(fname, inp, model, atol=atol, is_fp8=True, te_outputs=te_outputs)

@skip_FP8
@pytest.mark.parametrize("scale_factor", [448])
@pytest.mark.parametrize(
    "precision,             atol", [
    [torch.float32,         1e-5],
    [torch.float16,         1e-5],
    [torch.bfloat16,        5e-3],
    ["fake-torch.bfloat16", 5e-3]
])
def test_export_gelu_fp8(scale_factor: float, precision: torch.dtype, atol: float):
    fake_bf16_io = precision == "fake-torch.bfloat16"
    # reset precision to torch.bfloat16 after capturing fake BF16 mode
    precision = torch.bfloat16 if precision == "fake-torch.bfloat16" else precision

    class TestFP8_Gelu(nn.Module):
        def __init__(self, fake_bf16_io):
            super().__init__()
            self.fp8_tensor = 0
            self.meta = create_meta(scale_factor)
            self.highprec_type = as_te_type(precision)
            self.fp8_type = tex.DType.kFloat8E4M3
            self.fake_bf16_io = fake_bf16_io

        def forward(self, inp):
            ret = gelu(
                inp,
                self.meta,
                self.fp8_tensor,
                self.fp8_type)
            ret = cast_from_fp8(
                ret,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                self.highprec_type)
            if self.fake_bf16_io:
                ret = ret.type(torch.float32)
            return ret

    # Set dimensions (these are arbitrary).
    in_features = 64
    hidden_size = 256
    inp = torch.randn(hidden_size, in_features, device="cuda",
        dtype=torch.float if fake_bf16_io else precision)
    high_prec_str = dtype2str(precision, fake_bf16_io=fake_bf16_io)
    fname = f"te.gelu_fp8_{scale_factor}{high_prec_str}.onnx"
    model = TestFP8_Gelu(fake_bf16_io)
    do_export(model, inp, fname)
    te_outputs = te_infer(model, inp, is_fp8=True)
    serialize_inputs_outputs(fname, inp, te_outputs)
    if fake_bf16_io or precision != torch.bfloat16:
        validate_result(fname, inp, model, rtol=0, atol=atol, is_fp8=True, allow_cnt_errors=2, te_outputs=te_outputs)

@pytest.mark.parametrize("scale_factors",
    [(224, 224,),
])
@pytest.mark.parametrize(
    "precision,             use_fp8, use_bias, use_gelu", [
    (torch.float32,         False,   False,    False),
    (torch.float16,         False,   False,    False),
    (torch.bfloat16,        False,   False,    False),
    (torch.float32,         False,   True,     False),
    (torch.float16,         False,   True,     False),
    (torch.bfloat16,        False,   True,     False),
    (torch.float32,         False,   True,     True),
    (torch.float16,         False,   True,     True),
    (torch.bfloat16,        False,   True,     True),

    # For FP8 GEMM GeLU is not used.
    (torch.float32,         True,    False,    False),
    (torch.float16,         True,    False,    False),
    (torch.bfloat16,        True,    False,    False),
    # When enabling bias we must use float16 or bfloat16 (because of kernel limitations)
    (torch.float16,         True,    True,     False),
    (torch.bfloat16,        True,    True,     False),
])
def test_export_gemm(
    seed_default_rng,
    precision, # Precision of inputs, weights, output and bias
    use_fp8,
    use_bias,
    use_gelu,
    scale_factors
):
    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    class TestFP8_GEMM(nn.Module):
        def __init__(self, precision, use_bias, gelu, scale_factors):
            super().__init__()
            self.use_bias = use_bias
            self.gelu = gelu
            self.precision = precision

            self.fp8_tensor_inp = tex.FP8FwdTensors.GEMM1_INPUT
            self.fp8_tensor_weight = tex.FP8FwdTensors.GEMM1_WEIGHT
            nb_inp_scales, nb_weight_scales = 1, out_features
            act_scale_factor, weight_scale_factor = scale_factors
            self.meta_inp = create_meta(act_scale_factor, nb_inp_scales)
            self.meta_weight = create_meta(weight_scale_factor, nb_weight_scales)

            bias_size = nb_weight_scales
            self.bias = torch.randn(bias_size, dtype=precision, device="cuda")
            self.gelu_input = torch.randn(hidden_size, out_features, dtype=precision, device="cuda")

            self.inp_type = tex.DType.kFloat8E4M3
            self.weights_type = tex.DType.kFloat8E4M3
            self.outp_type = precision

        def forward(self, inp, weight):
            inp_fp8 = cast_to_fp8(
                inp,
                self.meta_inp,
                self.fp8_tensor_inp,
                self.inp_type)

            weight_fp8 = cast_to_fp8(
                weight,
                self.meta_weight,
                self.fp8_tensor_weight,
                self.weights_type)

            ret, _ = fp8_gemm(
                weight_fp8,
                self.meta_weight.scale_inv,
                self.fp8_tensor_weight,
                self.inp_type,
                inp_fp8,
                self.meta_inp.scale_inv,
                self.fp8_tensor_inp,
                self.weights_type,
                self.outp_type,
                get_workspace(),
                bias=self.bias,
                use_bias=self.use_bias,
                use_split_accumulator=False)
            return ret

    class Test_GEMM(nn.Module):
        def __init__(self, precision, use_bias=False, gelu=False):
            super().__init__()
            self.use_bias = use_bias
            self.gelu = gelu
            self.precision = precision
            bias_size = out_features
            self.bias = torch.randn(bias_size, dtype=precision, device="cuda")
            self.gelu_input = torch.randn(hidden_size, out_features, dtype=precision, device="cuda")

        def forward(self, inp, weight):
            outp_type = self.precision

            # note: due to logic in lines 104:116 and L129 in cpp_extensions.py
            # it appears either bias OR gelu can be activated, not both
            ret, _, _ = gemm(
                weight,
                inp,
                outp_type,
                get_workspace(),

                # test bias
                bias=self.bias,
                use_bias=self.use_bias,

                # test gelu
                gelu=self.gelu,
                gelu_input=self.gelu_input,
                grad=False, # only True for backward pass
                accumulate=False,
            )
            return ret

    # If gelu is applied then bias must be added, as defined by TE kernel.
    if use_gelu: assert use_bias
    # Set dimensions (these are arbitrary).
    out_features = 128
    hidden_size = 256
    in_features = 64
    inp = torch.randn(hidden_size, in_features, device="cuda", dtype=precision)
    weight = torch.randn(out_features, in_features, device="cuda", dtype=precision)
    fp8_str = "_fp8" if use_fp8 else ""
    bias_str = "_bias" if use_bias else ""
    gelu_str = "_gelu" if use_gelu else ""
    high_prec_str = dtype2str(precision)
    fname = f"te.gemm{fp8_str}{bias_str}{gelu_str}{high_prec_str}.onnx"
    input_names = ['input', 'weight']
    if use_fp8:
        model = TestFP8_GEMM(precision, use_bias, use_gelu, scale_factors)
        do_export(model, (inp, weight), fname, use_fp8, input_names=input_names)
        te_outputs = te_infer(model, (inp, weight), is_fp8=use_fp8)
        serialize_inputs_outputs(fname, (inp, weight), te_outputs, input_names=input_names)
        if precision != torch.bfloat16:
            validate_result(fname, (inp, weight), model, rtol=1e-2, atol=2e-2,
                is_fp8=True, input_names=input_names, te_outputs=te_outputs)
    else:
        model = Test_GEMM(precision, use_bias, use_gelu)
        do_export(model, (inp, weight), fname, use_fp8, input_names=input_names)
        te_outputs = te_infer(model, (inp, weight), is_fp8=use_fp8)
        serialize_inputs_outputs(fname, (inp, weight), te_outputs, input_names=input_names)
        if precision != torch.bfloat16:
            validate_result(fname, (inp, weight), model, rtol=1e-2, atol=2e-2,
                input_names=input_names, te_outputs=te_outputs)

@pytest.mark.parametrize("scale_factor", [448, 112])
@pytest.mark.parametrize("zero_centered_gamma", [False, True])
@pytest.mark.parametrize(
    "use_fp8, precision,             atol", [
    [False,   torch.float32,         1e-7],
    [False,   torch.float16,         1e-7],
    [False,   torch.bfloat16,        1e-7],
    [False,   "fake-torch.bfloat16", 1e-7],
    [True,    torch.float32,         1e-7],
    [True,    torch.float16,         1e-7],
    [True,    torch.bfloat16,        1e-2],
    [True,    "fake-torch.bfloat16", 1e-2]
])
def test_export_layernorm(
    seed_default_rng,
    use_fp8: bool,
    scale_factor: float,
    precision: torch.dtype,
    zero_centered_gamma: bool,
    atol: float
):
    fake_bf16_io = precision == "fake-torch.bfloat16"
    # reset precision to torch.bfloat16 after capturing fake BF16 mode
    precision = torch.bfloat16 if precision == "fake-torch.bfloat16" else precision

    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    # Set dimensions (these are arbitrary).
    inp_shape = [64, 32]

    class Test_Layernorm(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            eps = 1e-6 # An arbitrary small value
            dtype = torch.float if fake_bf16_io else precision
            self.ln = te.LayerNorm(inp_shape[1], eps, params_dtype=dtype,
                zero_centered_gamma=False).eval().cuda()

        def forward(self, inp):
            ret = self.ln(inp)
            return ret

    class TestFP8_Layernorm(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            normalized_shape = torch.Size(inp.shape[1:])
            self.weight = torch.randn(*normalized_shape, device="cuda",
                dtype=torch.float32 if fake_bf16_io else precision)
            self.bias = torch.zeros(*normalized_shape, device="cuda",
                dtype=torch.float32 if fake_bf16_io else precision)
            self.eps = 1e-6 # An arbitrary small value

            self.fp8_tensor = tex.FP8FwdTensors.GEMM1_INPUT
            self.meta = create_meta(scale_factor)
            self.fp8_type = tex.DType.kFloat8E4M3

        def forward(self, inp):
            ret = texcpp.layernorm_fwd_fp8_inf(
                inp,
                self.weight,
                self.bias,
                self.eps,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                zero_centered_gamma)

            ret = cast_from_fp8(
                ret,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                as_te_type(precision))
            if fake_bf16_io:
                ret = ret.type(torch.float32)
            return ret

    inp = torch.randn(*inp_shape, device="cuda", dtype=torch.float32 if fake_bf16_io else precision)
    model = TestFP8_Layernorm() if use_fp8 else Test_Layernorm()
    high_prec_str = dtype2str(precision, fake_bf16_io=fake_bf16_io)
    fp8_str = f"_fp8-{scale_factor}" if use_fp8 else ""
    fname = f"te.layernorm{fp8_str}{high_prec_str}.onnx"
    do_export(model, inp, fname, use_fp8=use_fp8)
    te_outputs = te_infer(model, inp, is_fp8=use_fp8)
    serialize_inputs_outputs(fname, inp, te_outputs)
    if fake_bf16_io or precision != torch.bfloat16:
        validate_result(
            fname, inp, model, atol=atol, is_fp8=use_fp8, allow_cnt_errors=3, te_outputs=te_outputs)

@pytest.mark.parametrize("scale_factor", [448, 112])
@pytest.mark.parametrize(
    "use_fp8, precision,             atol", [
    [False,   torch.float32,         1e-7],
    [False,   torch.float16,         1e-7],
    [False,   torch.bfloat16,        1e-7],
    [False,   "fake-torch.bfloat16", 1e-7],
    [True,    torch.float32,         1e-7],
    [True,    torch.float16,         1e-7],
    [True,    torch.bfloat16,        1e-2],
    [True,    "fake-torch.bfloat16", 1e-2]
])
def test_export_rmsnorm(
    seed_default_rng,
    use_fp8: bool,
    scale_factor: float,
    precision: torch.dtype,
    atol: float
):
    fake_bf16_io = precision == "fake-torch.bfloat16"
    # reset precision to torch.bfloat16 after capturing fake BF16 mode
    precision = torch.bfloat16 if precision == "fake-torch.bfloat16" else precision

    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    # Set dimensions (these are arbitrary).
    inp_shape = [64, 32]

    class Test_RMSnorm(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            eps = 1e-6 # An arbitrary small value
            dtype = torch.float if fake_bf16_io else precision
            self.ln = te.RMSNorm(inp_shape[1], eps, params_dtype=dtype).eval().cuda()

        def forward(self, inp):
            ret = self.ln(inp)
            return ret

    class TestFP8_RMSnorm(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            normalized_shape = torch.Size(inp.shape[1:])
            self.weight = torch.randn(*normalized_shape, device="cuda",
                dtype=torch.float32 if fake_bf16_io else precision)
            self.eps = 1e-6 # An arbitrary small value

            self.fp8_tensor = tex.FP8FwdTensors.GEMM1_INPUT
            self.meta = create_meta(scale_factor)
            self.fp8_type = tex.DType.kFloat8E4M3

        def forward(self, inp):
            ret = texcpp.rmsnorm_fwd_fp8_inf(
                inp,
                self.weight,
                self.eps,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                False)

            ret = cast_from_fp8(
                ret,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                as_te_type(precision))
            if fake_bf16_io:
                ret = ret.type(torch.float32)
            return ret

    inp = torch.randn(*inp_shape, device="cuda", dtype=torch.float32 if fake_bf16_io else precision)
    model = TestFP8_RMSnorm() if use_fp8 else Test_RMSnorm()
    high_prec_str = dtype2str(precision, fake_bf16_io=fake_bf16_io)
    fp8_str = f"_fp8-{scale_factor}" if use_fp8 else ""
    fname = f"te.layernorm{fp8_str}{high_prec_str}.onnx"
    do_export(model, inp, fname, use_fp8=use_fp8)
    te_outputs = te_infer(model, inp, is_fp8=use_fp8)
    serialize_inputs_outputs(fname, inp, te_outputs)
    if fake_bf16_io or precision != torch.bfloat16:
        validate_result(
            fname, inp, model, atol=atol, is_fp8=use_fp8, allow_cnt_errors=3, te_outputs=te_outputs)

@pytest.mark.parametrize("scale_factor", [1])
@pytest.mark.parametrize("use_fp8", [False, True])
# Returning the bias is a TE fusion optimization we don't care about.
@pytest.mark.parametrize("return_bias", [False])
@pytest.mark.parametrize(
    "precision,      use_bias",[
    (torch.float32,  False),
    (torch.float32,  True),
    (torch.float16,  False),
    (torch.float16,  True),
    # Todo: cannot configure BF16 when bias is disabled (ORT issue?)
    (torch.bfloat16, False),
    # Todo: cannot configure BF16 when bias is enabled (ORT issue?)
    (torch.bfloat16, True),
])
def test_export_linear(
    seed_default_rng,
    scale_factor: float,
    use_fp8: bool,
    use_bias: bool,
    return_bias: bool,
    precision: torch.dtype
):
    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    # Set dimensions (these are arbitrary).
    in_features = 64
    out_features = 256
    hidden_size = 256

    class Test_Linear(nn.Module):
        def __init__(self,
                in_features,
                out_features,
                use_bias,
                return_bias,
                precision
            ):
            super().__init__()
            self.linear = te.Linear(
                in_features,
                out_features,
                bias=use_bias,
                return_bias=return_bias,
                params_dtype=precision
            )

        def forward(self, inp):
            ret = self.linear(inp)
            return ret

    inp = torch.randn(hidden_size, in_features, device="cuda", dtype=precision)
    fp8_str = "_fp8" if use_fp8 else ""
    bias_str = "_bias" if use_bias else ""
    high_prec_str = dtype2str(precision)
    fname = f"te.linear{fp8_str}{bias_str}{high_prec_str}.onnx"
    with te.fp8_autocast(enabled=use_fp8):
        model = Test_Linear(
            in_features,
            out_features,
            use_bias,
            return_bias,
            precision
        ).to(device='cuda')
        if use_fp8:
            set_layer_scale(model.linear, scale_factor, num_gemms=1)
        do_export(model, inp, fname, use_fp8)
        te_outputs = te_infer(model, inp, is_fp8=use_fp8)
        serialize_inputs_outputs(fname, inp, te_outputs)

        if precision in (torch.bfloat16, ):
            return
        if not use_fp8:
            validate_result(fname, inp, model, atol=1e-3, te_outputs=te_outputs)
        else:
            validate_result(fname, inp, model, atol=1e-3, is_fp8=use_fp8, te_outputs=te_outputs)

@pytest.mark.parametrize("scale_factor", [112])
@pytest.mark.parametrize("use_fp8", [False, True])
# Returning the bias is a TE fusion optimization we don't care about.
@pytest.mark.parametrize("return_bias", [False])
@pytest.mark.parametrize("return_layernorm_output", [False])
@pytest.mark.parametrize(
    "precision,      use_bias",[
    (torch.float32,  False),
    (torch.float32,  True),
    (torch.float16,  True),
    (torch.float16,  False),
    (torch.bfloat16, True),
    (torch.bfloat16, False),
])
@pytest.mark.parametrize("zero_centered_gamma", [False, True])
@pytest.mark.parametrize("normalization", all_normalizations)
def test_export_layernorm_linear(
    seed_default_rng,
    scale_factor: float,
    use_fp8: bool,
    use_bias: bool,
    return_bias: bool,
    return_layernorm_output: bool,
    precision: torch.dtype,
    zero_centered_gamma: bool,
    normalization: str,
):
    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    if normalization == "RMSNorm" and zero_centered_gamma:
        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")

    # Set dimensions (these are arbitrary).
    in_features = 64
    out_features = 256
    hidden_size = 256

    inp = torch.randn(in_features, out_features, device="cuda", dtype=precision)
    fp8_str = "_fp8" if use_fp8 else ""
    bias_str = "_bias" if use_bias else ""
    high_prec_str = dtype2str(precision)
    fname = f"te.layernorm_linear{fp8_str}{bias_str}{high_prec_str}.onnx"

    with te.fp8_autocast(enabled=use_fp8):
        model = te.LayerNormLinear(
            hidden_size,
            3 * hidden_size,
            bias=use_bias,
            return_bias=return_bias,
            return_layernorm_output=return_layernorm_output,
            params_dtype=precision,
            zero_centered_gamma=zero_centered_gamma,
            normalization=normalization,
        ).to(device='cuda')
        if use_fp8:
            set_layer_scale(model, scale_factor, num_gemms=1)
        do_export(model, inp, fname, use_fp8)
        te_outputs = te_infer(model, inp, is_fp8=use_fp8)
        serialize_inputs_outputs(fname, inp, te_outputs)
        if precision in (torch.bfloat16, ):
            return
        if not use_fp8:
            validate_result(fname, inp, model, atol=1e-3, te_outputs=te_outputs)
        elif precision != torch.bfloat16:
            validate_result(fname, inp, model, atol=1e-6, is_fp8=use_fp8, te_outputs=te_outputs)

@pytest.mark.parametrize("scale_factor", [112])
@pytest.mark.parametrize("use_fp8", [False, True])
# Returning the bias is a TE fusion optimization we don't care about.
@pytest.mark.parametrize("return_bias", [False])
@pytest.mark.parametrize("return_layernorm_output", [False])
@pytest.mark.parametrize(
    "precision,      use_bias",[
    (torch.float32,  False),
    (torch.float32,  True),
    (torch.float16,  True),
    (torch.float16,  False),
    (torch.bfloat16, True),
    (torch.bfloat16, False),
])
@pytest.mark.parametrize("zero_centered_gamma", [False, True])
@pytest.mark.parametrize("activation", supported_activations)
@pytest.mark.parametrize("normalization", all_normalizations)
def test_export_layernorm_mlp(
    seed_default_rng,
    scale_factor: float,
    use_fp8: bool,
    use_bias: bool,
    return_bias: bool,
    return_layernorm_output: bool,
    precision: torch.dtype,
    zero_centered_gamma: bool,
    activation: str,
    normalization: str,
):
    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    if normalization == "RMSNorm" and zero_centered_gamma:
        pytest.skip("RMSNorm does not support zero_centered_gamma yet!")

    # Set dimensions (these are arbitrary).
    in_features = 64
    out_features = 256
    hidden_size = 256
    ffn_hidden_size = 256

    inp = torch.randn(in_features, out_features, device="cuda", dtype=precision)
    fp8_str = "_fp8" if use_fp8 else ""
    bias_str = "_bias" if use_bias else ""
    high_prec_str = dtype2str(precision)
    fname = f"te.layernorm_mlp{fp8_str}{bias_str}{high_prec_str}_{activation}.onnx"
    with te.fp8_autocast(enabled=use_fp8):
        model = te.LayerNormMLP(
            hidden_size,
            ffn_hidden_size,
            bias=use_bias,
            return_bias=return_bias,
            return_layernorm_output=return_layernorm_output,
            params_dtype=precision,
            zero_centered_gamma=zero_centered_gamma,
            activation=activation,
            normalization=normalization,
        ).to(device='cuda')
        if use_fp8:
            set_layer_scale(model, scale_factor, num_gemms=2)
        do_export(model, inp, fname, use_fp8)
        te_outputs = te_infer(model, inp, is_fp8=use_fp8)
        serialize_inputs_outputs(fname, inp, te_outputs)
        if precision in (torch.bfloat16, ):
            return
        atol = 1e-6 if use_fp8 else (5e-1 if activation=="swiglu" else 1e-3)
        validate_result(fname, inp, model, atol=atol, is_fp8=use_fp8, te_outputs=te_outputs)

@skip_FP8
@pytest.mark.parametrize(
    "precision,      use_mask, attn_mask_type", [
    (torch.float32,  True,     "arbitrary"), # calls forward_torch_softmax (apply user mask)
    (torch.float32,  False,    "no_mask"),   # calls forward_torch_softmax (apply no mask)
    (torch.float16,  False,    "causal"),    # calls forward_torch_softmax (apply dynamic onnx mask)
    (torch.float16,  True,     "arbitrary"), # calls forward_torch_softmax (apply user mask)
    (torch.float16,  False,    "no_mask"),   # calls forward_torch_softmax (apply no mask)
    (torch.bfloat16, False,    "causal"),    # calls forward_torch_softmax (apply dynamic onnx mask)
    (torch.bfloat16, True,     "arbitrary"), # calls forward_torch_softmax (apply user mask)
    (torch.bfloat16, False,    "no_mask"),   # calls forward_torch_softmax (apply no mask)
])
def test_export_core_attention(
    seed_default_rng,
    set_max_seq_len,
    precision: torch.dtype,
    use_mask: bool,
    attn_mask_type: str,
):
    # Set dimensions (these are arbitrary).
    seq_len, batch_size, num_attention_heads, kv_channels = (64, 4, 1, 64)
    qkv_size = (seq_len, batch_size, num_attention_heads, kv_channels)
    qkv_format = "sbhd"

    query_layer = torch.randn(qkv_size, dtype=precision, device="cuda")
    key_layer = torch.randn(qkv_size, dtype=precision, device="cuda")
    value_layer = torch.randn(qkv_size, dtype=precision, device="cuda")
    input_names = ["query", "key", "value", "attention_mask"]
    attention_mask = None
    if use_mask:
        # Generate a random mask with 50% probability for 0 or 1.
        probs = 0.5 * torch.ones(batch_size, 1, 1, seq_len, device="cuda", dtype=precision)
        attention_mask = torch.bernoulli(probs).to("cuda", dtype=torch.bool)
    inp = (query_layer, key_layer, value_layer, attention_mask)

    mask_str = get_attn_mask_str(use_mask, attn_mask_type)
    high_prec_str = dtype2str(precision)
    fname = f"te.core_attention{mask_str}{high_prec_str}.onnx"

    model = te.attention.DotProductAttention(
        num_attention_heads=num_attention_heads,
        kv_channels=kv_channels,
        attention_dropout=0.5,
        qkv_format=qkv_format,
        attn_mask_type=attn_mask_type,
    ).to(device='cuda')
    do_export(model,
            inp,
            fname,
            input_names=input_names,
            use_fp8=True)
    te_outputs = te_infer(model, inp, is_fp8=True)
    serialize_inputs_outputs(fname, inp, te_outputs, input_names=input_names)
    if precision in (torch.bfloat16, ):
        return
    validate_result(fname, inp, model, is_fp8=True, atol=1e-2, input_names=input_names, te_outputs=te_outputs)

test_configs_multihead_attention = [
    #"use_mask, attn_mask_type"
    (False,    "no_mask"),   # calls ScaledSoftmax
    (True,     "arbitrary"), # calls ScaledMaskedSoftmax
]
test_configs_attention_type = [
    #"input_layernorm, attention_type, fuse_qkv_params"
    (True,             "self",         True),
    (False,            "self",         True),
    (True,             "self",         False),
    (False,            "self",         False),
    (True,             "cross",        True),
    (False,            "cross",        True),
    (True,             "cross",        False),
    (False,            "cross",        False),
]
@pytest.mark.parametrize("use_fp8", [False, True])
@pytest.mark.parametrize("use_mask, attn_mask_type", test_configs_multihead_attention)
@pytest.mark.parametrize("precision", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("return_layernorm_output", [False])
@pytest.mark.parametrize("input_layernorm, attention_type, fuse_qkv_params", test_configs_attention_type)
def test_export_multihead_attention(
    seed_default_rng,
    set_max_seq_len,
    use_fp8: bool,
    use_mask: bool,
    attn_mask_type: str,
    precision: torch.dtype,
    return_layernorm_output: bool,
    input_layernorm: bool,
    attention_type: str,
    fuse_qkv_params: bool
):
    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    hidden_size = 256
    sequence_length = 128
    batch_size = 4
    num_attention_heads = 32
    kv_channels = 8
    attention_dropout = 0.1
    layernorm_epsilon = 1e-5
    init_method = output_layer_init_method = get_default_init_method()
    attention_args = (
        hidden_size,
        num_attention_heads,
        kv_channels,
        attention_dropout,
        layernorm_epsilon,
        init_method,
        output_layer_init_method,
    )

    hidden_states_context = torch.randn(sequence_length, batch_size, hidden_size, dtype=precision, device="cuda")
    attention_mask = None
    if use_mask and attn_mask_type != "causal":
        # Generate a random mask with 50% probability for 0 or 1.
        probs = 0.5 * torch.ones(batch_size, 1, sequence_length, sequence_length, device="cuda", dtype=precision)
        attention_mask = torch.bernoulli(probs).to("cuda", dtype=torch.bool)

    encoder_output = None

    if attention_type == "cross":
        encoder_output = torch.randn(sequence_length, batch_size, hidden_size, dtype=precision, device="cuda")

    fp8_str = "_fp8" if use_fp8 else ""
    dtype_str = dtype2str(precision)
    attn_type_str = "_self-attention" if attention_type == "self" else "_cross-attention"
    fuse_qkv_str = "_fused-qkv" if fuse_qkv_params else ""
    attn_mask_str = get_attn_mask_str(use_mask, attn_mask_type)
    input_ln_str = "_input-ln" if input_layernorm else ""
    fname = f"te.multihead_attention{fp8_str}{attn_mask_str}{attn_type_str}{input_ln_str}{fuse_qkv_str}{dtype_str}.onnx"

    model = te.MultiheadAttention(
        *attention_args,
        attn_mask_type=attn_mask_type,
        params_dtype=precision,
        return_layernorm_output=return_layernorm_output,
        input_layernorm=input_layernorm,
        attention_type=attention_type,
        fuse_qkv_params=fuse_qkv_params,
        return_bias=True,
    ).to(device='cuda')

    inp_context = (hidden_states_context, attention_mask, encoder_output)
    input_names = ["hidden_states", "attention_mask", "encoder_output"]
    output_names=["attention_output", "attention_bias"]
    do_export(model, inp_context, fname, use_fp8, input_names=input_names, output_names=output_names,
        dynamic_axes={"hidden_states": {0: "seq", 1:"bs"},
                      "attention_output": {0: "seq", 1:"bs"}})
    te_outputs = te_infer(model, inp_context, is_fp8=use_fp8)
    serialize_inputs_outputs(fname, inp_context, te_outputs, input_names=input_names, output_names=output_names)
    if precision in (torch.bfloat16, ):
        return

    if not use_fp8:
        validate_result(fname, inp_context, model, atol=1e-3, input_names=input_names,
            output_names=output_names, te_outputs=te_outputs)
    else:
        validate_result(fname, inp_context, model, atol=1e-2, is_fp8=use_fp8,
            input_names=input_names, output_names=output_names, allow_cnt_errors=3,
            te_outputs=te_outputs)

    # In GPT generative phase (inference) the input sequence is smaller than the maximum
    # allowed sequence length and we want to test this condition.
    # Pretend that we're in generative phase when it makes sense (causal mask and self-attention).
    is_generative_phase = (attn_mask_type == "causal" and attention_type == "self")
    if is_generative_phase:
        seq_len_offset = 8
        hidden_states_generative = torch.randn(sequence_length-seq_len_offset, batch_size, hidden_size, dtype=precision, device="cuda")
        inp_generative = (hidden_states_generative, attention_mask, encoder_output)
        if not use_fp8:
            validate_result(fname, inp_generative, model, atol=1e-3, input_names=input_names, output_names=output_names)
        else:
            validate_result(fname, inp_generative, model, atol=1e-2, is_fp8=use_fp8,
                input_names=input_names, output_names=output_names, allow_cnt_errors=3)

@pytest.mark.parametrize("use_fp8", [False, True])
@pytest.mark.parametrize("use_mask, attn_mask_type", test_configs_multihead_attention)
@pytest.mark.parametrize("output_layernorm", [
    #True, # TO DO: handle this
    False
])
@pytest.mark.parametrize("precision", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("fuse_qkv_params", [False, True])
@pytest.mark.parametrize("zero_centered_gamma", [False, True])
@pytest.mark.parametrize("activation", supported_activations)
def test_export_transformer_layer(
    seed_default_rng,
    set_max_seq_len,
    use_fp8: bool,
    use_mask: bool,
    attn_mask_type: str,
    output_layernorm: bool,
    precision: torch.dtype,
    fuse_qkv_params: bool,
    zero_centered_gamma: bool,
    activation: str,
):
    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    # Layer configuration
    hidden_size = 64
    sequence_length = 128
    batch_size = 1
    ffn_hidden_size = 256
    num_attention_heads = 4

    input_tensor = torch.rand(sequence_length, batch_size, hidden_size, dtype=precision, device="cuda")
    input_names = ["input", "attention_mask"]
    attention_mask = None
    if use_mask and attn_mask_type != "causal":
        # Generate a random mask with 50% probability for 0 or 1.
        probs = 0.5 * torch.ones(batch_size, 1, sequence_length, sequence_length, device="cuda", dtype=precision)
        attention_mask = torch.bernoulli(probs).to("cuda", dtype=torch.bool)
    inp = (input_tensor, attention_mask)

    fp8_str = "_fp8" if use_fp8 else ""
    fuse_qkv_params_str = "_fused-qkv" if fuse_qkv_params else ""
    high_prec_str = dtype2str(precision)
    attn_mask_str = get_attn_mask_str(use_mask, attn_mask_type)
    fname = f"te.transformer_layer{fp8_str}{attn_mask_str}{fuse_qkv_params_str}{high_prec_str}_{activation}.onnx"

    model = te.TransformerLayer(
        hidden_size,
        ffn_hidden_size,
        num_attention_heads,
        self_attn_mask_type=attn_mask_type,
        output_layernorm=output_layernorm,
        params_dtype=precision,
        fuse_qkv_params=fuse_qkv_params,
        zero_centered_gamma=zero_centered_gamma,
        activation=activation).to(device='cuda')
    do_export(model, inp, fname, use_fp8, input_names=input_names)
    te_outputs = te_infer(model, inp, is_fp8=use_fp8)
    serialize_inputs_outputs(fname, inp, te_outputs, input_names=input_names)
    if precision in (torch.bfloat16, ):
        return
    atol = 5e-1 if use_fp8 else (5e-1 if activation=="swiglu" else 1e-3)
    validate_result(fname, inp, model, atol=atol, is_fp8=use_fp8, input_names=input_names, te_outputs=te_outputs)

@pytest.mark.parametrize("use_fp8", [True])
@pytest.mark.parametrize("ln_scale_factor", [448*2])
@pytest.mark.parametrize("gemm_scale_factors", [(224, 224,),])
@pytest.mark.parametrize("precision", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("zero_centered_gamma", [False, True])
def test_export_gemm_layernorm(
    seed_default_rng,
    use_fp8: bool,
    ln_scale_factor: float,
    gemm_scale_factors: Tuple[float, float],
    precision: torch.dtype,
    zero_centered_gamma: bool
):
    """This is a regression test for testing that all LN inputs have the same type.

    The test sets up GEMM with FP32 output which feeds into an LN that is configured
    with FP16 or BF16 weights and bias.
    """

    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)
    class TestFP8_GemmLayernorm(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            normalized_shape = torch.Size(inp.shape[1:])
            self.weight = torch.randn(*normalized_shape, dtype=precision, device="cuda")
            self.bias = torch.zeros(*normalized_shape, dtype=precision, device="cuda")
            self.eps = 1e-6 # An arbitrary small value

            self.fp8_tensor = tex.FP8FwdTensors.GEMM1_INPUT
            self.meta = create_meta(ln_scale_factor)
            self.fp8_type = tex.DType.kFloat8E4M3
            self.gemm = TestFP8_GEMM(
                precision, use_bias=False, gelu=False, scale_factors=gemm_scale_factors)

        def forward(self, inp, weight):
            x = self.gemm(inp, weight)
            x = texcpp.layernorm_fwd_fp8_inf(
                x,
                self.weight,
                self.bias,
                self.eps,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                zero_centered_gamma)

            x = cast_from_fp8(
                x,
                self.meta,
                self.fp8_tensor,
                self.fp8_type,
                tex.DType.kFloat32 if precision == torch.float32 else tex.DType.kFloat16)
            return x

    out_features = 128
    hidden_size = 128
    in_features = 128
    class TestFP8_GEMM(nn.Module):
        def __init__(self, precision, use_bias, gelu, scale_factors):
            super().__init__()
            self.use_bias = use_bias
            self.gelu = gelu
            self.precision = precision

            self.fp8_tensor_inp = tex.FP8FwdTensors.GEMM1_INPUT
            self.fp8_tensor_weight = tex.FP8FwdTensors.GEMM1_WEIGHT
            nb_inp_scales, nb_weight_scales = 1, out_features
            act_scale_factor, weight_scale_factor = scale_factors
            self.meta_inp = create_meta(act_scale_factor, nb_inp_scales)
            self.meta_weight = create_meta(weight_scale_factor, nb_weight_scales)

            bias_size = nb_weight_scales
            self.bias = torch.randn(bias_size, dtype=precision, device="cuda")
            self.gelu_input = torch.randn(hidden_size, out_features, dtype=precision, device="cuda")

            self.inp_type = tex.DType.kFloat8E4M3
            self.weights_type = tex.DType.kFloat8E4M3
            self.outp_type = precision

        def forward(self, inp, weight):
            inp_fp8 = cast_to_fp8(
                inp,
                self.meta_inp,
                self.fp8_tensor_inp,
                self.inp_type)

            weight_fp8 = cast_to_fp8(
                weight,
                self.meta_weight,
                self.fp8_tensor_weight,
                self.weights_type)

            ret, _ = fp8_gemm(
                weight_fp8,
                self.meta_weight.scale_inv,
                self.fp8_tensor_weight,
                self.inp_type,
                inp_fp8,
                self.meta_inp.scale_inv,
                self.fp8_tensor_inp,
                self.weights_type,
                self.outp_type,
                get_workspace(),
                bias=self.bias,
                use_bias=self.use_bias,
                use_split_accumulator=False)
            return ret

    inp = torch.randn(hidden_size, in_features, dtype=precision, device="cuda")
    weight = torch.randn(out_features, in_features, dtype=precision, device="cuda")
    model = TestFP8_GemmLayernorm()
    high_prec_str = dtype2str(precision)
    fp8_str = f"_fp8" if use_fp8 else ""
    fname = f"te.gemm_layernorm{fp8_str}{high_prec_str}.onnx"
    input_names = ['input', 'weight']
    do_export(model, (inp, weight), fname, use_fp8=use_fp8, input_names=input_names)
    te_outputs = te_infer(model, (inp, weight), is_fp8=use_fp8)
    serialize_inputs_outputs(fname, (inp, weight), te_outputs, input_names=input_names)
    if precision not in (torch.bfloat16, ):
        validate_result(
            fname, (inp, weight), model, atol=5e-2, is_fp8=use_fp8, allow_cnt_errors=2,
            input_names=input_names, te_outputs=te_outputs)

@skip_FP8
@pytest.mark.parametrize("use_fp8", [True, False])
@pytest.mark.parametrize("precision", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("zero_centered_gamma", [True])
def test_export_gpt_generation(
    seed_default_rng,
    set_max_seq_len,
    use_fp8: bool,
    precision: torch.dtype,
    zero_centered_gamma: bool,
):
    """Test that the ONNX model can correctly handle inputs with different shapes and that
    the attention mask it adjusted on-the-fly to different sequence lengths.
    """

    # Skip FP8 tests on non-hopper devices
    if use_fp8 and not fp8_available:
        pytest.skip(reason_for_no_fp8)

    # Layer configuration
    hidden_size = 64
    sequence_length = 128
    batch_size = 1
    ffn_hidden_size = 256
    num_attention_heads = 4
    attention_mask = None
    use_mask = True
    attn_mask_type = "causal"
    fuse_qkv_params = True
    output_layernorm = False

    fp8_str = "_fp8" if use_fp8 else ""
    fuse_qkv_params_str = "_fused-qkv" if fuse_qkv_params else ""
    high_prec_str = dtype2str(precision)
    attn_mask_str = get_attn_mask_str(use_mask, attn_mask_type)
    fname = f"te.transformer_layer_generative{fp8_str}{attn_mask_str}{fuse_qkv_params_str}{high_prec_str}.onnx"

    model = te.TransformerLayer(
        hidden_size,
        ffn_hidden_size,
        num_attention_heads,
        self_attn_mask_type=attn_mask_type,
        output_layernorm=output_layernorm,
        params_dtype=precision,
        fuse_qkv_params=fuse_qkv_params,
        zero_centered_gamma=zero_centered_gamma).to(device='cuda')

    # "Context phase": use full input sequence length
    input_names = ["input"]
    output_names = ["output"]
    input_tensor = torch.rand(sequence_length, batch_size, hidden_size, dtype=precision, device="cuda")
    inp = (input_tensor,)
    do_export(model, inp, fname, use_fp8,
        input_names=input_names, output_names=output_names,
        dynamic_axes={"input": {0: "seq", 1:"bs"},
                      "output": {0: "seq", 1:"bs"}, })
    te_outputs = te_infer(model, inp, is_fp8=use_fp8)
    serialize_inputs_outputs(fname, inp, te_outputs, input_names=input_names, output_names=output_names)
    if precision not in (torch.bfloat16, ):
        validate_result(fname, inp, model, atol=6e-3, is_fp8=use_fp8, input_names=input_names,
            te_outputs=te_outputs)

    # "Generative phase": use a single input (sequence len=1). For FP8 we need to pad the sequence to mult of 8.
    sequence_length = 1 if not use_fp8 else 8
    input_tensor = torch.rand(sequence_length, batch_size, hidden_size, dtype=precision, device="cuda")
    inp = (input_tensor, attention_mask)
    te_outputs = te_infer(model, inp, is_fp8=use_fp8)
    serialize_inputs_outputs(fname, inp, te_outputs, input_names=input_names)
    if precision not in (torch.bfloat16, ):
        validate_result(fname, inp, model, atol=6e-3, is_fp8=use_fp8, input_names=input_names,
            te_outputs=te_outputs)

@pytest.mark.parametrize("enabled", [True, False])
def test_export_ctx_manager(enabled):
    assert is_in_onnx_export_mode() == False
    with te.onnx_export(enabled):
        assert is_in_onnx_export_mode() == enabled
    assert is_in_onnx_export_mode() == False


import news_parser as Util
import datetime
import time
import DBHandler
from tensorflow.keras.models import load_model
import predict
CompanyList=[]
Headless = True # False : 창띄움, True : 창없음
MakeCompanyList = False # 회사 리스트 갱신
host = '데이터베이스 주소'
ID= '계정명'
PW='비밀번호'
DB_name='DB이름'

def GetNewsInfo(driver):
    headlines, news_info, Text,NewsUrl = Util.GetNews(driver) # 헤드라인, 신문사 정보 및 게시 시간, 본문, 기사링크 파싱.
    CompanyFromNews = Util.GetCompanyFromNews(headlines, CompanyList)
    #Util.save_headlines(headlines, news_info, Text,CompanyFromNews,NewsUrl)
    #Util.PrintNews(headlines, news_info, Text, CompanyFromNews)
    return headlines, news_info, Text, NewsUrl,CompanyFromNews

def GetPriceInfo(driver):
    NameList, PriceInfo, Fluctuation = Util.get_prices(driver) #KTOP 30, KOSPI, KOSPI200, KOSDAQ, KOSDAQ150, KRX300 순
    Util.PrintPrice(NameList, PriceInfo, Fluctuation)
    return NameList, PriceInfo, Fluctuation

def MakeCompanyFile(MakeCompanyList):
    #Company CSV파일 생성
    Util.MakeCompanyCSV()

if __name__ == '__main__':
    print("Setting Interface...")
    CompanyList = Util.GetCompanyList() # 코스피 상장 기업 업로드
    try: #네이버 증권
        NewsDriver = Util.News_get_driver(Headless)
    except Exception as ex:
        print("News Driver Err")
        print('에러가 발생 했습니다', ex)
    try : #한국거래소
         PriceDriver = Util.NowPriceDriver(Headless)
    except Exception as ex:
         print("Price Driver Err")
         print('에러가 발생 했습니다', ex)
    try: # 다음 증권
        KospiImageDriver = Util.Get_KospiGraphDriver(Headless)
    except Exception as ex:
        print("KospiImage Driver Err")
        print('에러가 발생 했습니다', ex)
    MakeCompanyFile(MakeCompanyList) #기업 리스트 갱신
    DBController = DBHandler.MySqlController(host, ID, PW, DB_name)

    label=[]
    while(True):
            now = datetime.datetime.now()
            nowDatetime = now.strftime('%Y_%m_%d_%H시%M분%S초'.encode('unicode-escape').decode()).encode().decode('unicode-escape')
            nowDatehour = now.strftime('%Y_%m_%d_%H시%M분'.encode('unicode-escape').decode()).encode().decode('unicode-escape')
            try:
                NameList, PriceInfo, Fluctuation = Util.get_prices(PriceDriver)
                Util.PrintPrice(NameList,PriceInfo,Fluctuation)
                DBController.update_totalprice(PriceInfo,Fluctuation)
            except Exception as ex:
                print("Price Info Err")
                print('에러가 발생 했습니다', ex)
                PriceDriver.quit()
                PriceDriver = Util.NowPriceDriver(Headless)
                NameList, PriceInfo, Fluctuation = Util.get_prices(PriceDriver)
            # print("========================================")
            try:
                headlines, news_info, Text,NewsUrl,CompanyFromNews = GetNewsInfo(NewsDriver) #뉴스에서 기업 추출
                print("News Updated...")
            except Exception as ex:
                print("News Update Err")
                NewsDriver.quit()
                NewsDriver = Util.News_get_driver(Headless)
                print('에러가 발생 했습니다', ex)
            try:
                Util.Write_News(headlines, CompanyFromNews, nowDatehour)  # 기업별 뉴스 자료 Writing
            except Exception as ex:
                print("News Write Err")
                CompanyList = Util.GetCompanyList()  # 코스피 상장 기업 업로드
            try:
                Util.GetKospiGraph(KospiImageDriver, PriceInfo, Fluctuation) # Kospi, Kosdaq 그래프 이미지 저장
                print("Get Kospi Graph")
            except Exception as ex:
                KospiImageDriver.quit()
                KospiImageDriver = Util.Get_KospiGraphDriver(Headless)
                print("Graph Err")
                print('에러가 발생 했습니다', ex)
            try:
                label = predict.classification(headlines, model)
                print("Get labels")
                DBController.UpdateNews(CompanyFromNews, headlines, Text, NewsUrl, news_info, label)  # 최신 20개 기사 DB저장
                DBController.InsertNewsHistory(CompanyFromNews, headlines, Text, NewsUrl, news_info, nowDatehour)
                print("DB Commit : News Updated, News History Inserted")
            except Exception as ex:
                print("Label Err")
                MakeCompanyFile(MakeCompanyList)  # 기업 리스트 갱신
                DBController = DBHandler.MySqlController(host, ID, PW, DB_name)
                print('에러가 발생 했습니다', ex)
            time.sleep(30)
            NewsDriver.refresh()
            PriceDriver.refresh()
            KospiImageDriver.refresh()
            print("DONE")


import os, sys, random
from datetime import datetime
platforms = {
    "darwin": "MacOS",
    "win32": "Windows32",
    "win64": "Windows64",
    "linux": "Linux",
}

def run_tool():
    while True:
        now = datetime.now()
        current_time = now.strftime("%H:%M:%S")
        command = input(f"[{current_time}] > ")
        if command == "os":
            print(f"Detected operating system: {platforms.get(sys.platform, 'Unknown OS')}")

        elif command == "clear":
            os.system("clear")

        elif "+" in command:
            nums = command.split("+")
            result = 0
            for x in nums:
                result += int(x)
            print(result)

        elif command == "games":
            print("""1. Guess the word
                     2. Rock paper scissor
                     3. Mad libs generator""")
            which = input("\nWhich game would you like to play? ")
            if which == "1":
                try:
                    with open("words.txt", "r") as file:
                        words = [word.strip() for word in file]
                    word_to_guess = random.choice(words)
                    guess = input(f"What word has a length of {len(word_to_guess)} characters, starts with {word_to_guess[0]} and ends with {word_to_guess[-1]}? ")
                    if guess != word_to_guess:
                        print(f"Incorrect. The word was {word_to_guess}.")

                    else:
                        print("Correct!")

                except FileNotFoundError:
                    print("Can't find words.txt. Make sure it's in the same directory and not named something else.")

            elif which == "2":
                choices = ["rock", "paper", "scissor"]
                user_choice = input("Enter your choice: ").lower()
                computer_choice = random.choice(choices)
                if user_choice == computer_choice:
                    print("Tie.")
                elif user_choice == "rock" and computer_choice == "scissor":
                    print("Rock beats scissor. You win!")

                elif user_choice == "rock" and computer_choice == "paper":
                    print("Paper beats rock. Computer wins!")

                elif user_choice == "scissor" and computer_choice == "rock":
                    print("Rock beats scissor. Computer wins!")

                elif user_choice == "scissor" and computer_choice == "paper":
                    print("Scissor beats paper. You win!")

                elif user_choice == "paper" and computer_choice == "scissor":
                    print("Scissor beats paper. Computer wins!")

                elif user_choice == "paper" and computer_choice == "rock":
                    print("Paper beats rock. You win!")

                else:
                    print("Invalid choice.")

            elif which == "3":
                give_me_words = input("Enter 5 words: ").split()
                print(f"The {give_me_words[0]} was slimy, it tasted like {give_me_words[1]}. I stayed at {give_me_words[2]} place. He made me {give_me_words[3]} in the morning. We packed up at {give_me_words[4]}.")

        elif command == "quit":
            sys.exit()

        else:
            print(f'Unknown command: "{command}" ')

run_tool()


"""
Defines a function wrapper that can be used to group calls to a 'listable'
function into batches.
"""

import sys
import traceback
import asyncio

from fsc.export import export

from . import wrap_to_coroutine

@export
class BatchSubmitter:
    """
    Function wrapper that collects calls to a function of one parameter, and submits
    it in batches to a function which can take a list of parameters.

    Arguments
    ---------
    func: Callable
        Function or coroutine which is "listable", i.e. given a list of input
        parameters it will return a list of results.
    loop: EventLoop
        The event loop on which the batch submitter runs. Uses
        ``asyncio.get_event_loop()`` by default.
    timeout: float
        Maximum time after which the batch submitter will submit all current
        tasks, even if the minimum batch size is not reached.
    sleep_time : float
        Time the batch submitter will sleep between checking if the minimum
        batch size has been reached.
    wait_batch_size : int
        Minimum batch size that will be submitted before the timeout has been
        reached. Is set to the same value as ``max_batch_size`` unless
        specified explicitly.
    max_batch_size : int
        The maximum size of a batch that will be submitted.
    """

    def __init__(
        self,
        func,
        *,
        loop=None,
        timeout=0.1,
        sleep_time=0.,
        wait_batch_size=None,
        max_batch_size=1000
    ):
        self._func = wrap_to_coroutine(func)
        self._loop = loop or asyncio.get_event_loop()
        self._timeout = timeout
        self._sleep_time = sleep_time

        if max_batch_size <= 0:
            raise ValueError('max_batch_size must be positive')
        self._max_batch_size = max_batch_size
        if wait_batch_size is None:
            wait_batch_size = self._max_batch_size
        if wait_batch_size <= 0:
            raise ValueError('wait_batch_size must be positive')
        self._wait_batch_size = wait_batch_size

        self._tasks = asyncio.Queue()
        self._batches = dict()
        self._submit_loop_task = None
        self._last_call_time = None

    async def __call__(self, x):
        """
        Adds a task for the given input, and starts the submission loop if needed.
        """
        fut = self._loop.create_future()
        self._tasks.put_nowait((x, fut))
        self._last_call_time = self._loop.time()
        if self._submit_loop_task is None or self._submit_loop_task.done():
            self._submit_loop_task = asyncio.Task(
                self._submit_loop(), loop=self._loop
            )
            self._submit_loop_task.add_done_callback(self._abort_on_exception)
        return await fut

    async def _submit_loop(self):
        """
        Waits for tasks and then creates the batches which evaluate the function.
        """
        while self._tasks.qsize() > 0:
            await self._wait_for_tasks()
            self._launch_batch()

    @staticmethod
    def _abort_on_exception(fut):
        """
        Callback that forces a SystemExit when there is an exception in the submit loop.
        """
        try:
            fut.result()
        except Exception:  # pylint: disable=broad-except
            sys.exit(''.join(traceback.format_exception(*sys.exc_info())))

    async def _wait_for_tasks(self):
        """
        Waits until either the timeout has passed or the queue size is big enough.
        """
        assert self._tasks.qsize() > 0
        while self._loop.time() - self._last_call_time < self._timeout:
            if self._tasks.qsize() >= self._wait_batch_size:
                return
            await asyncio.sleep(self._sleep_time)

    def _launch_batch(self):
        """
        Launch a calculation batch.
        """
        inputs = []
        futures = []
        for _ in range(self._max_batch_size):
            try:
                key, fut = self._tasks.get_nowait()
                inputs.append(key)
                futures.append(fut)
            except asyncio.QueueEmpty:
                break
        task = asyncio.ensure_future(self._func(inputs))
        task.add_done_callback(self._process_finished_batch)
        self._batches[task] = futures

    def _process_finished_batch(self, batch_future):
        """
        Assign the results / exceptions to the futures of all finished batches.
        """
        task_futures = self._batches.pop(batch_future)
        try:
            results = batch_future.result()
            assert len(results) == len(task_futures)
            for fut, res in zip(task_futures, results):
                fut.set_result(res)
        except Exception as exc:  # pylint: disable=broad-except
            for fut in task_futures:
                fut.set_exception(exc)


#Usage: python3 kptable-appendix-11b.py [-h] [--help]
import datetime
import pathlib
import pandas as pd
import xlsxwriter
from lukeghg.crf.crfxmlconstants import ch4co2eq, n2oco2eq, ctoco2
from lukeghg.crf.crfxmlfunctions import ConvertFloat,ConvertSign, ConvertToCO2, SumTwoValues, SumBiomassLs
from lukeghg.crf.crfxmlfunctions import PaddingList, GenerateRowTitleList

#These constants will come from CrfXMLConstants
#Please check the
#ch4co2eq=25.0
#n2oco2eq=298.0
#ctoco2=44.0/12.0
#nkset={'IE','NA','NO','NE'}
#Sort the Records based on YearName in ascending order
def SortKey(x):
    return x.attrib['YearName']
#---------------------------------The main program begins--------------------------------------------------

def appendix11b(start,end,directory,file_name):
    #Command line generator
    global ch4co2eq, n2oco2eq,ctoco2
    inventory_year=end
    #Output file, the table
    kptable_appendix_11b_file = file_name
    directory=directory+'/'
    #Table Appendix-11b Afforestation/Reforestation and Deforestation files
    #Deforestation "Conversion to water CH4" comes from NIR folder
    kp4a2_fl_to_waters_ch4_org=directory+'KP4A2_FLtowaters_orgsoil_CH4.csv'
    #The rest of the files are from  crf-folder.
    kp4a_agr_bm_gains_losses=directory+'KP4A_agr_bm_gains_losses.csv'
    #2015 file name change
    #2016 file name change
    #2015:kpa2_ar_under_D_gains='KP4A2_AR_und_defor_treebm_gains.csv'
    kpa2_ar_under_D_gains=directory+'KP4A2_AR_und_D_living_biomass_gains_trees.csv'
    kp4a2_sl_soil=directory+'KP4A2_SL_soil.csv'
    kp4a2_ar_under_d_soil=directory+'KP4A2_AR_und_defor_soils.csv'
    #2015 KP_MTT_UID.csv in two files: KP_defor_mineral.csv and KP_defor_organic.csv
    #kp_uid_mtt='KP_MTT_UID.csv'
    kp_defor_mineral=directory+'KP_defor_mineral.csv'
    kp_defor_organic=directory+'KP_defor_organic.csv'
    kp4a2_fl_to_wl_soil=directory+'KP4A2_FLtoWL_soils.csv'
    kp4a2_clglpewesl_deadwood=directory+'KP4A2_CLGLPEWESL_deadwood.csv'
    kp4a2_d_living_biomass_losses_trees=directory+'KP4A2_D_living_biomass_losses_trees.csv'
    kp4a2_fl_to_waters_org_soil=directory+'KP4A2_FLtowaters_orgsoil.csv'
    #2015 rename
    #kp4a2_d_mineralization='KP4A2_D_mineraalisationcl_gl_sl.csv'
    kp4a2_d_mineralization=directory+'KPA2_soil_leaching_N2O.csv'
    #2015 addition is Afforestation mineralization
    kp4_ar_mineralization=directory+'KP4_Affor_mineralization.csv'
    kp4a2_fl_to_wl_non_co2=directory+'KP4A2_FLtoWL_soils_nonCO2.csv'
    kp4_living_biomass_gains_trees=directory+'KP4_living_biomass_gains_trees.csv'
    kp4_ar_living_biomass_losses_trees=directory+'KP4A1_AR_living_biomass_losses_trees.csv'
    kp4a1_clglsl_mineral_soil=directory+'KP4A1_CLGLSL_mineral_soil.csv'
    kp4a1_ar_org_soil=directory+'KP4A1_AR_Organic_soil_C.csv'
    kp4a11_wildfires=directory+'KP4A11_wildfires.csv'
    kp4a1_clglpewesl_organic_soils_nonco2=directory+'KP4A1_CLGLPEWESL_organic_soils_nonco2.csv'
    kp4_hwp_ard=directory+'KP4_HWP-AR.csv'
    #Data for the two Tables in Appendix 1
    #1. Deforestation: Conversion to water CH4
    #Change in 2015: use third (CH4) line in kp4a2_fl_to_waters_org_soil
    #----------------------------------------
    f = open(kp4a2_fl_to_waters_org_soil)
    #Read to a list [[year,val1],[year,val2]....,[year,valN]]
    ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #Third line is CH4
    ls = ls[2]
    #Convert to CO2
    ls.pop(0)
    fl_to_waters_ch4_co2_ls = [ConvertToCO2(ch4co2eq,x) for x in ls]
    f.close()
    #2. Agriculture Afforestation and Deforestation biomasses
    #---------------------------------------------------------
    f = open(kp4a_agr_bm_gains_losses)
    agr_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #Pick the deforestation, the first part in the file
    agr_d_ls = agr_ls[0:8]
    #Pick the Afforestation and Reforestation, rest of the file
    agr_ar_ls = agr_ls[8:len(agr_ls)]
    #Deforestation: Sum the biomassess: Cropland, Grassland, North and South Finland,
    #above ground and below ground
    agr_d_bm_ls=SumBiomassLs(agr_d_ls)
    #Afforestation: Sum the biomassess: Cropland, Grassland, North and South Finland,
    #above ground and below ground
    agr_ar_bm_ls=SumBiomassLs(agr_ar_ls)
    #Convert to CO2 and convert sign: if biomass increases -> emissiosn decrease and vice versa
    agr_d_co2_ls=[ConvertSign(ConvertToCO2(ctoco2,x)) for x in agr_d_bm_ls]
    agr_ar_co2_ls=[ConvertSign(ConvertToCO2(ctoco2,x)) for x in agr_ar_bm_ls]
    f.close()
    #3. Deforestation biomass losses in trees
    #----------------------------------------
    f = open(kp4a2_d_living_biomass_losses_trees)
    d_trees_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #Sum the biomasses, CL, GL, SETT, WL, North and South Finland, below and above ground
    trees_bm_ls = SumBiomassLs(d_trees_ls)
    #Convert to CO2 and convert sign: if biomass increases -> emissions decrease and vice versa
    trees_d_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in trees_bm_ls]
    f.close()
    #Deforestation Biomass: Afforestation/Reforestation under Deforestation, gains
    #---------------------------------------------------------------------
    f = open(kpa2_ar_under_D_gains)
    ar_under_d_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    ar_under_d_sum_ls = SumBiomassLs(ar_under_d_ls)
    #Convert to CO2 and convert sign: if biomass increases -> emissions decrease and vice versa
    ar_under_d_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in ar_under_d_sum_ls]
    f.close()
    #Deforestation: DOM+SOM Mineral soils
    #-----------------------------------
    f = open(kp4a2_sl_soil)
    d_sl_soil_ls =  [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_sl_soil_sum_ls = SumBiomassLs(d_sl_soil_ls)
    d_sl_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_sl_soil_sum_ls]
    f.close()
    f = open(kp4a2_ar_under_d_soil)
    d_ar_under_d_soil_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #Sum the both lines
    d_ar_under_d_min_soil_sum_ls = SumBiomassLs(d_ar_under_d_soil_ls)
    d_ar_under_d_min_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_ar_under_d_min_soil_sum_ls]
    f.close()
    #Settlements are now in Mineral soil, take lines 9 and 10
    f = open(kp4a2_clglpewesl_deadwood)
    d_clglpewesl_deadwood_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_clglpewesl_deadwood_ls = d_clglpewesl_deadwood_ls[8:len(d_clglpewesl_deadwood_ls)]
    d_clglpewesl_deadwood_min_sum_ls = SumBiomassLs(d_clglpewesl_deadwood_ls)
    d_clglpewesl_deadwood_min_co2_ls =  [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_clglpewesl_deadwood_min_sum_ls]
    f.close()
    f = open(kp_defor_mineral)
    d_mtt_min_soil_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_mtt_min_soil_ls = d_mtt_min_soil_ls[0:4]
    d_mtt_min_soil_sum_ls = SumBiomassLs(d_mtt_min_soil_ls)
    d_mtt_min_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_mtt_min_soil_sum_ls]
    f.close()
    d_dom_som_min_soil_deadwood_sum_co2_ls = [SumTwoValues(a,SumTwoValues(b,SumTwoValues(c,d))) for (a,b,c,d) in zip(d_sl_soil_co2_ls,
                                                                                d_ar_under_d_min_soil_co2_ls,
                                                                                d_clglpewesl_deadwood_min_co2_ls,
                                                                                d_mtt_min_soil_co2_ls)]

    #Deforestation: DOM+SOM Organic soils + Deadwood
    #-----------------------------------------------
    f = open(kp4a2_fl_to_wl_soil)
    d_fl_to_wl_soil_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_fl_to_wl_soil_sum_ls = SumBiomassLs(d_fl_to_wl_soil_ls)
    d_fl_to_wl_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_fl_to_wl_soil_sum_ls]
    f.close()
    f = open(kp4a2_clglpewesl_deadwood)
    d_clglpewesl_deadwood_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_clglpewesl_deadwood_ls = d_clglpewesl_deadwood_ls[0:8]
    d_clglpewesl_deadwood_org_sum_ls = SumBiomassLs(d_clglpewesl_deadwood_ls)
    d_clglpewesl_deadwood_org_co2_ls =  [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_clglpewesl_deadwood_org_sum_ls]
    f.close()
    f = open(kp_defor_organic)
    d_mtt_org_soil_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_mtt_org_soil_sum_ls = SumBiomassLs(d_mtt_org_soil_ls)
    d_mtt_org_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_mtt_org_soil_sum_ls]
    #print(8,d_mtt_org_soil_co2_ls)
    f.close()
    d_dom_som_org_soil_deadwood_sum_co2_ls=[SumTwoValues(a,SumTwoValues(b,c)) for (a,b,c) in
                                            zip(d_fl_to_wl_soil_co2_ls,d_clglpewesl_deadwood_org_co2_ls,d_mtt_org_soil_co2_ls)]
    #Deforestation: Conversion to water CO2
    #Change in 2015: Lines 1 and 2 are C
    #--------------------------------------
    f = open(kp4a2_fl_to_waters_org_soil)
    d_fl_to_waters_org_soil_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #Lines 1,2 are C
    d_fl_to_waters_org_soil_ls = d_fl_to_waters_org_soil_ls[0:2]
    d_fl_to_waters_org_soil_sum_ls = SumBiomassLs(d_fl_to_waters_org_soil_ls)
    d_fl_to_waters_org_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in d_fl_to_waters_org_soil_sum_ls]
    f.close()
    #Deforestation: Mineralization
    #-----------------------------
    f = open(kp4a2_d_mineralization)
    d_mineralization_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #2015 two last lines are N2O
    d_mineralization_ls = d_mineralization_ls[2:len(d_mineralization_ls)]
    d_mineralization_sum_ls =  SumBiomassLs(d_mineralization_ls)
    d_mineralization_co2_ls = [ConvertToCO2(n2oco2eq,x) for x in d_mineralization_sum_ls]
    f.close()
    #Deforestation: Drained and rewetted organic soils N2O
    #-----------------------------------------------------
    f = open(kp4a2_fl_to_wl_non_co2)
    d_fl_to_wl_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_fl_to_wl_n2o_ls = d_fl_to_wl_ls[0:1]
    d_fl_to_wl_n2o_sum_ls =  SumBiomassLs(d_fl_to_wl_n2o_ls)
    d_fl_to_wl_n2o_co2_ls = [ConvertToCO2(n2oco2eq,x) for x in d_fl_to_wl_n2o_sum_ls]
    f.close()
    #Deforestation: Drained and rewetted organic soils CH4
    #-----------------------------------------------------
    f = open(kp4a2_fl_to_wl_non_co2)
    d_fl_to_wl_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    d_fl_to_wl_ch4_ls = d_fl_to_wl_ls[1:2]
    d_fl_to_wl_ch4_sum_ls =  SumBiomassLs(d_fl_to_wl_ch4_ls)
    d_fl_to_wl_ch4_co2_ls = [ConvertToCO2(ch4co2eq,x) for x in d_fl_to_wl_ch4_sum_ls]
    f.close()
    #Deforestation HWP
    #-----------------
    #HWP is IE
    d_hwp_ls = ['IE']*(inventory_year-1990+1)
    #4. Afforestation living biomass gains and losses trees
    #-------------------------------------------
    f = open(kp4_living_biomass_gains_trees)
    ar_trees_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    f.close()
    #Pick the Afforestation part, 2015 mineral and organic are added (not separately in the file)
    ar_bm_gains_trees_ls = ar_trees_ls[4:len(ar_trees_ls)]
    #Sum the biomasses, CL, CL, WLpeat, WLorg, Settlement, mineral, orgaing, South and North Finland
    ar_bm_sum_gains_trees_ls = SumBiomassLs(ar_bm_gains_trees_ls)
    f = open(kp4_ar_living_biomass_losses_trees)
    ar_bm_losses_trees_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    f.close()
    ar_bm_sum_losses_trees_ls = SumBiomassLs(ar_bm_losses_trees_ls)
    ar_bm_net_trees_ls = [SumTwoValues(x,y) for (x,y) in zip(ar_bm_sum_gains_trees_ls,ar_bm_sum_losses_trees_ls)]
    #Convert to CO2 and convert sign: if biomass increases -> emissions decrease and vice versa
    trees_ar_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in ar_bm_net_trees_ls]
    #5. Afforestation, DOM+SOM Mineral soils
    #----------------------------------------
    f = open(kp4a1_clglsl_mineral_soil)
    dom_som_min_soil_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    dom_som_min_soil_sum_ls = SumBiomassLs(dom_som_min_soil_ls)
    #Convert to CO2 and convert sign: if biomass increases -> emissions decrease and vice versa
    dom_som_min_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in dom_som_min_soil_sum_ls]
    f.close()
    #6. Afforestation, DOM+SOM Organinc soils
    #----------------------------------------
    f=open(kp4a1_ar_org_soil)
    dom_som_org_soil_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    dom_som_org_soil_sum_ls = SumBiomassLs(dom_som_org_soil_ls)
    #Convert to CO2 and convert sign: if biomass increases -> emissions decrease and vice versa
    dom_som_org_soil_co2_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in dom_som_org_soil_sum_ls ]
    f.close()
    #7. Biomass burning
    #------------------
    f=open(kp4a11_wildfires)
    biomass_burning_ls = [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #CO2 South and North Finland
    bm_burning_co2_ls = biomass_burning_ls[0:2]
    bm_burning_co2_sum_ls = SumBiomassLs(bm_burning_co2_ls)
    #CH4 South and North Finland
    bm_burning_ch4_ls = biomass_burning_ls[2:4]
    bm_burning_ch4_sum_ls = SumBiomassLs(bm_burning_ch4_ls)
    #N2O South and North Finland
    bm_burning_n2o_ls  = biomass_burning_ls[4:6]
    bm_burning_n2o_sum_ls = SumBiomassLs(bm_burning_n2o_ls)
    #Convert ch4 and n2o to co2eq and sum all three emissions
    biomass_burning_ch4co2eq_ls = [ConvertToCO2(ch4co2eq,x) for x in bm_burning_ch4_sum_ls]
    biomass_burning_n2oco2eq_ls = [ConvertToCO2(n2oco2eq,x) for x in bm_burning_n2o_sum_ls]
    biomass_burning_co2_sum_ls = [SumTwoValues(x,SumTwoValues(y,z)) for (x,y,z) in zip(bm_burning_co2_sum_ls,biomass_burning_ch4co2eq_ls,biomass_burning_n2oco2eq_ls)]
    #print(biomass_burning_co2_sum_ls)
    f.close()
    #8. 2015 addition Mineralization
    #-------------------------------
    f=open(kp4_ar_mineralization)
    ar_mineralization_ls=[x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #South and North Fianland
    ar_mineralization_no2_ls = ar_mineralization_ls[0:2]
    ar_mineralization_n2o_sum_ls=SumBiomassLs(ar_mineralization_no2_ls)
    ar_mineralization_n2o_co2_ls = [ConvertToCO2(n2oco2eq,x) for x in ar_mineralization_n2o_sum_ls]
    #9.Drained organic soils N2O
    #---------------------------
    f=open(kp4a1_clglpewesl_organic_soils_nonco2)
    drained_org_soils_nonco2_ls =  [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #Two lines in the file, the first one is CH4
    drained_org_soils_sum_ch4_ls = drained_org_soils_nonco2_ls[0:1]
    drained_org_soils_sum_ch4_ls =  SumBiomassLs(drained_org_soils_sum_ch4_ls)
    #Convert from N2O to CO2
    drained_org_soils_sum_ch4co2eq_ls = [ConvertToCO2(ch4co2eq,x) for x in drained_org_soils_sum_ch4_ls]
    f.close()
    #10.Drained organic soils CH4
    #---------------------------
    f=open(kp4a1_clglpewesl_organic_soils_nonco2)
    drained_org_soils_nonco2_ls =  [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #Two lines in the file, the second one is CH4
    drained_org_soils_sum_n2o_ls = drained_org_soils_nonco2_ls[1:2]
    drained_org_soils_sum_n2o_ls = SumBiomassLs(drained_org_soils_sum_n2o_ls)
    #Convert from CH4 to CO2
    drained_org_soils_sum_n2oco2eq_ls = [ConvertToCO2(n2oco2eq,x) for x in drained_org_soils_sum_n2o_ls]
    f.close()
    #11.HWP afforestation
    #--------------------
    f=open(kp4_hwp_ard)
    hwp_ard_ls =  [x.rpartition('#')[2].split() for x in f.readlines() if x.count('#') != 1]
    #2015 the file structure is in 3 parts: 1) Initial stock, 2) Gains and losses 3) half-time information
    #The sum of gains and losses will go to the table
    #2016 the file structure is for each item: 1) half life, 2) init stock, 3) gains and 4) losses
    #Pick every fourth item starting from the first right place in the list
    hwp_ar_ls_gains = hwp_ard_ls[2::4]
    hwp_ar_ls_losses = hwp_ard_ls[3::4]
    hwp_ar_ls = hwp_ar_ls_gains+hwp_ar_ls_losses
    #print(hwp_ar_ls)
    hwp_ar_sum_ls = SumBiomassLs(hwp_ar_ls)
    #Removals are good for the atmosphere, change the sign
    hwp_ar_sum_ls = [ConvertSign(ConvertToCO2(ctoco2,x)) for x in hwp_ar_sum_ls]
    #HWP does not have full time series from 1990
    hwp_padding_ls = PaddingList(inventory_year,hwp_ar_sum_ls)
    hwp_padding_ls = ['IE']*len(hwp_padding_ls)
    hwp_ar_sum_co2_ls = hwp_padding_ls+hwp_ar_sum_ls
    f.close()

    #Create the two tables Afforestation/Reforestation and Deforestation
    #-------------------------------------------------------------------
    print("Creating first text file for Afforestation/Reforestation and Deforestation in", file_name)
    f1 = open(kptable_appendix_11b_file,'w')
    delim ='#'
    table_name="Appendix_11b"
    table_header="Net emissions and removals from the ativities under Articles 3.3\n"
    table1title1="Table 1_App_11b Net emissions and removals from Afforestation and Reforestation, kt CO2eq.\n"
    table2title2="Table 2_App_11b Net emissions and removals from Deforestation, ktCO2eq.\n"
    table1columns1=delim+"Biomass"+delim+"DOM+SOM Mineral soils"+delim+"DOM+SOM Organic soils"+delim+"Biomass burning"+delim+"Mineralization"+delim
    table1columns1=table1columns1+"Drained organic soils N2O"+delim+"Drained organic soils CH4"+delim+"HWP"+delim+"Total\n"
    table2columns2=delim+"Biomass"+delim+"DOM+SOM Mineral soils"+delim+"DOM+SOM Organic soils+Deadwood"+delim+"Conversion to water CO2"+delim+"Mineralization"+delim
    table2columns2=table2columns2+"Drained and rewetted organic soils CH4"+delim+"Drained and rewetted organic soils NO2"+delim
    table2columns2=table2columns2+"HWP"+delim+"Conversion to water CH4"+delim+"Total"+"#\n"
    #Row titles from 199 to inventory year
    row_title_ls = GenerateRowTitleList(start,inventory_year)
    f1.write(table_name)
    f1.write(table_header)
    #Afforestation and Reforestation
    f1.write(table1title1)
    f1.write(table1columns1)
    for (year,agr_ar_co2,trees_ar_co2,dom_som_min,dom_som_org,bm_burning,ar_min_co2,n2oco2eq,ch4co2eq,hwp) in zip(row_title_ls,agr_ar_co2_ls,trees_ar_co2_ls,
                                                                                dom_som_min_soil_co2_ls,
                                                                                dom_som_org_soil_co2_ls,
                                                                                biomass_burning_co2_sum_ls,
                                                                                ar_mineralization_n2o_co2_ls,
                                                                                drained_org_soils_sum_n2oco2eq_ls,
                                                                                drained_org_soils_sum_ch4co2eq_ls,
                                                                                hwp_ar_sum_co2_ls):
        total=SumTwoValues(SumTwoValues(SumTwoValues(SumTwoValues(SumTwoValues(SumTwoValues(SumTwoValues(SumTwoValues(agr_ar_co2,trees_ar_co2),
                                                                                dom_som_min),dom_som_org),bm_burning),
                                                                                ar_min_co2),n2oco2eq),ch4co2eq),hwp)
        f1.write(str(year)+delim+str(SumTwoValues(agr_ar_co2,trees_ar_co2))+"#"+str(dom_som_min)+"#"+str(dom_som_org)+"#"+str(bm_burning)+"#")
        f1.write(str(ar_min_co2)+"#"+str(n2oco2eq)+"#"+str(ch4co2eq)+"#"+str(hwp)+"#"+str(total)+"#\n")
    f1.write("Data from:"+"#"+kp4a_agr_bm_gains_losses+" Lines:9-"+str(len(agr_ls))+"#"+kp4a1_clglsl_mineral_soil+"#"+kp4a1_ar_org_soil+"#"+kp4a11_wildfires+"#")
    f1.write(kp4_ar_mineralization+" Lines:1,2"+delim+kp4a1_clglpewesl_organic_soils_nonco2+" Line:1"+"#"+kp4a1_clglpewesl_organic_soils_nonco2+" Line:2"+"#"+kp4_hwp_ard+"#\n")
    f1.write("#"+kp4_living_biomass_gains_trees+" Lines:5-"+str(len(ar_trees_ls))+"#"+"#"+"#"+"CO2 Lines:1-2,CH4 Lines:3-4,N2O Lines:5-6"+"####"+"Gains Lines:3,7,11 etc."+"#\n")
    f1.write("#"+kp4_ar_living_biomass_losses_trees+"#######"+"Losses Lines:4,8,12 etc."+"#\n")
    f1.write('\n\n')

    #Deforestation
    f1.write(table2title2)
    f1.write(table2columns2)
    for (year,agr_d_co2,trees_d_co2,ar_under_d_co2,
        d_sl_soil_co2,d_ar_under_d_min_soil_co2,d_mtt_min_soil_co2,d_clglpewesl_deadwood_min_co2,
        dom_som_org_soil_deadwood_co2,
        d_fl_to_waters_org_soil_co2,
        d_mineralization_co2,
        d_fl_to_wl_ch4_co2,d_fl_to_wl_n2o_co2,
        d_hwp,
        to_waters_ch4) in zip(row_title_ls,agr_d_co2_ls,trees_d_co2_ls,ar_under_d_co2_ls,
                            d_sl_soil_co2_ls,d_ar_under_d_min_soil_co2_ls,d_mtt_min_soil_co2_ls,d_clglpewesl_deadwood_min_co2_ls,
                            d_dom_som_org_soil_deadwood_sum_co2_ls,
                            d_fl_to_waters_org_soil_co2_ls,
                            d_mineralization_co2_ls,
                            d_fl_to_wl_ch4_co2_ls,d_fl_to_wl_n2o_co2_ls,
                            d_hwp_ls,
                            fl_to_waters_ch4_co2_ls):
        biomass = SumTwoValues(agr_d_co2,trees_d_co2)
        biomass = SumTwoValues(biomass,ar_under_d_co2)
        dom_som_min_soil=SumTwoValues(d_sl_soil_co2,d_ar_under_d_min_soil_co2)
        dom_som_min_soil=SumTwoValues(dom_som_min_soil,d_mtt_min_soil_co2)
        dom_som_min_soil=SumTwoValues(dom_som_min_soil,d_clglpewesl_deadwood_min_co2)
        total1 = SumTwoValues(SumTwoValues(SumTwoValues(SumTwoValues(agr_d_co2,trees_d_co2),ar_under_d_co2),d_sl_soil_co2),d_ar_under_d_min_soil_co2)
        total2 = SumTwoValues(SumTwoValues(d_mtt_min_soil_co2,d_clglpewesl_deadwood_min_co2),dom_som_org_soil_deadwood_co2)
        total3 = SumTwoValues(SumTwoValues(d_fl_to_waters_org_soil_co2,d_mineralization_co2),d_fl_to_wl_ch4_co2)
        total4 = SumTwoValues(SumTwoValues(d_fl_to_wl_n2o_co2,d_hwp),to_waters_ch4)
        total = total1+total2+total3+total4
        f1.write(str(year)+delim+str(biomass)+delim+str(dom_som_min_soil)+delim+str(dom_som_org_soil_deadwood_co2)+delim+str(d_fl_to_waters_org_soil_co2)+delim+
                str(d_mineralization_co2)+delim+str(d_fl_to_wl_ch4_co2)+delim+str(d_fl_to_wl_n2o_co2)+delim+str(d_hwp)+delim+str(to_waters_ch4)+"#"+str(total)+"#\n")

    f1.write("Data from:"+"#"+kp4a_agr_bm_gains_losses+" Lines:1-8"+delim+kp4a2_sl_soil+delim+kp4a2_fl_to_wl_soil+delim+kp4a2_fl_to_waters_org_soil+" Lines:1,2"+"#")
    f1.write(kp4a2_d_mineralization+delim+kp4a2_fl_to_wl_non_co2+" Line:2"+delim+kp4a2_fl_to_wl_non_co2+" Line:1"+delim+"No file"+delim)
    f1.write(kp4a2_fl_to_waters_org_soil+" Line:3"+"#\n")
    f1.write("#"+kp4a2_d_living_biomass_losses_trees+delim+kp4a2_ar_under_d_soil+delim+kp4a2_clglpewesl_deadwood+" Lines:1-8"+"#\n")
    f1.write("#"+kpa2_ar_under_D_gains+delim+kp_defor_mineral+delim+kp_defor_organic+"#\n")
    f1.write("#"+delim+kp4a2_clglpewesl_deadwood+" Lines:9-10"+"#\n")
    now = datetime.datetime.now()
    #print(str(now))
    f1.write("Date produced: "+str(now)+"\n")
    f1.close()
    #Create excel
    p = pathlib.Path(file_name)
    stem = p.stem
    p_excel = pathlib.Path(stem+'.xlsx')
    print("Creating Excel file for Afforestation/Reforestation and Deforestation in", str(p_excel))
    #Define max number of columns, dataframe can adjust to it
    names=['col' + str(x) for x in range(12) ]
    df = pd.read_csv(file_name,engine='python',header=None,delimiter='#',keep_default_na=False,names=names,dtype=str)
    writer = pd.ExcelWriter(p_excel,engine='openpyxl')
    df_float = df.applymap(ConvertFloat)
    df_float.to_excel(writer,file_name,header=False)
    writer.close()


from dataclasses import dataclass
from datetime import datetime
from sqlalchemy_serializer import SerializerMixin
from app import db

@dataclass
class Result(db.Model, SerializerMixin):
    __tablename__ = 'results'
    name: str
    last_build: int
    last_result: str
    last_exception: str
    traceback: str
    url: str
    last_update: datetime

    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String())
    last_build = db.Column(db.Integer)
    last_result = db.Column(db.String())
    last_exception = db.Column(db.Text())
    traceback = db.Column(db.Text())
    url = db.Column(db.String())
    last_update = db.Column(db.DateTime, nullable=False,
                            default=datetime.utcnow)