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

"""TimesFM models in Flax."""

import dataclasses
import functools
import gc
import logging
import math
import os
from pathlib import Path
from typing import Any, Callable, Dict

import einshape
from flax import nnx
import huggingface_hub
import jax
import jax.numpy as jnp
import jaxtyping
import numpy as np
import orbax.checkpoint as ocp

from .. import configs
from ..flax import dense, transformer, util
from . import timesfm_2p5_base

jax_einshape = einshape.jax_einshape
scan = util.scan_along_axis
revin = util.revin

Float = jaxtyping.Float
Bool = jaxtyping.Bool
Array = jaxtyping.Array


def try_gc():
  for d in jax.local_devices():
    stats = d.memory_stats()
    if stats is None:
      return
    if stats["bytes_in_use"] / stats["bytes_limit"] > 0.75:
      gc.collect()
      break


@nnx.vmap(in_axes=(None, 0), out_axes=0)
def _create_stacked_transformers(
  config: configs.StackedTransformersConfig, key: jax.Array
):
  return transformer.Transformer(config.transformer, rngs=nnx.Rngs(key))


def _scan_along_axis(f, init, xs, axis: int, **kwargs):
  """Scans along an axis."""
  moved_xs = jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, axis, 0), xs)
  carry, moved_ys = jax.lax.scan(f, init, moved_xs, **kwargs)
  return (
    carry,
    jax.tree_util.tree_map(lambda x: jnp.moveaxis(x, 0, axis), moved_ys),
  )


@nnx.scan(in_axes=(0, nnx.Carry, None, 0), out_axes=(nnx.Carry, 0))
def _apply_stacked_transformers(
  model: transformer.Transformer,
  x: Float[Array, "b n d"],
  m: Float[Array, "b n"],
  decode_cache: util.DecodeCache | None = None,
) -> Float[Array, "b n d"]:
  return model(x, m, decode_cache=decode_cache)


class TimesFM_2p5_200M_flax_module(nnx.Module):  # pylint: disable=invalid-name
  """TimesFM 2.5 with 200M parameters."""

  config = timesfm_2p5_base.TimesFM_2p5_200M_Definition()
  decode_index: int = 5
  compiled_decode: Callable[..., Any] | None = None
  backend: str = ""
  context: int = 0
  horizon: int = 0
  per_core_batch_size: int = 0

  def __init__(self):
    super().__init__()
    self.backend = jax.devices()[0].platform
    self.num_devices = len(jax.devices(self.backend))

    # Names constants.
    self.p = self.config.input_patch_len  # 32
    self.o = self.config.output_patch_len  # 128
    self.os = self.config.output_quantile_len  # 1024
    self.m = self.o // self.p  # 4
    self.x = self.config.stacked_transformers.num_layers  # 20
    self.h = self.config.stacked_transformers.transformer.num_heads  # 16
    self.md = self.config.stacked_transformers.transformer.model_dims  # 1280
    self.hd = self.md // self.h  # 80
    self.q = len(self.config.quantiles) + 1  # 10
    self.aridx = self.config.decode_index  # 5

    # Layers.
    self.tokenizer = dense.ResidualBlock(self.config.tokenizer)
    self.stacked_xf = _create_stacked_transformers(
      self.config.stacked_transformers,
      jax.random.split(jax.random.key(42), self.x),
    )
    self.output_projection_point = dense.ResidualBlock(
      self.config.output_projection_point
    )
    self.output_projection_quantiles = dense.ResidualBlock(
      self.config.output_projection_quantiles
    )

  def __call__(
    self,
    inputs: Float[Array, "b n p"],
    masks: Bool[Array, "b n p"],
    decode_cache: util.DecodeCache | None = None,
  ):
    tokenizer_inputs = jnp.concatenate([inputs, masks.astype(inputs.dtype)], axis=-1)
    input_embeddings = self.tokenizer(tokenizer_inputs)
    if decode_cache is None:
      decode_cache = [None] * self.x
    output_embeddings, decode_cache = _apply_stacked_transformers(
      self.stacked_xf, input_embeddings, masks[..., -1], decode_cache
    )
    output_ts = self.output_projection_point(output_embeddings)
    output_quantile_spread = self.output_projection_quantiles(output_embeddings)
    return (
      input_embeddings,
      output_embeddings,
      output_ts,
      output_quantile_spread,
    ), decode_cache

  @nnx.jit(static_argnames=("horizon",))
  def decode(self, horizon: int, inputs, masks):
    batch_size, context = inputs.shape[0], inputs.shape[1]
    num_decode_steps = (horizon - 1) // self.o
    num_input_patches = context // self.p
    decode_cache_size = num_input_patches + num_decode_steps * self.m

    # Prefill
    patched_inputs = jax_einshape("b(np)->bnp", inputs, b=batch_size, p=self.p)
    patched_masks = jax_einshape("b(np)->bnp", masks, b=batch_size, p=self.p)
    (last_n, last_mu, last_sigma), (_, context_mu, context_sigma) = scan(
      lambda carry, xs: util.update_running_stats(*carry, *xs),
      init=(zero := jnp.zeros(shape=(batch_size)), zero, zero),
      xs=(patched_inputs, patched_masks),
      axis=1,
    )
    decode_cache = util.DecodeCache(
      next_index=jnp.zeros(shape=(self.x, batch_size), dtype=jnp.int32),
      num_masked=jnp.zeros(shape=(self.x, batch_size), dtype=jnp.int32),
      key=jnp.zeros(shape=(self.x, batch_size, decode_cache_size, self.h, self.hd)),
      value=jnp.zeros(shape=(self.x, batch_size, decode_cache_size, self.h, self.hd)),
    )
    normed_inputs = revin(patched_inputs, context_mu, context_sigma, reverse=False)
    normed_inputs = jnp.where(patched_masks, 0.0, normed_inputs)
    (_, _, normed_outputs, normed_quantile_spread), decode_cache = self(
      normed_inputs, patched_masks, decode_cache
    )
    renormed_outputs = jax_einshape(
      "bn(oq)->bnoq",
      revin(normed_outputs, context_mu, context_sigma, reverse=True),
      o=self.o,
      q=self.q,
    )
    renormed_quantile_spread = jax_einshape(
      "bn(oq)->bnoq",
      revin(normed_quantile_spread, context_mu, context_sigma, reverse=True),
      o=self.os,
      q=self.q,
    )[:, -1, ...]

    # Autogressive decode
    @nnx.scan(in_axes=(None, nnx.Carry, 0), out_axes=(nnx.Carry, 1))
    def _ar_decode(module, carry, unused_iter):
      last_renormed_output, (last_n, last_mu, last_sigma), decode_cache = carry
      new_patched_input = jax_einshape(
        "b(mp)->bmp", last_renormed_output, m=module.m, p=module.p
      )
      new_mask = jnp.zeros_like(new_patched_input, dtype=jnp.bool)
      carry_stats, (_, new_mu, new_sigma) = scan(
        lambda carry, xs: util.update_running_stats(*carry, *xs),
        init=(last_n, last_mu, last_sigma),
        xs=(new_patched_input, new_mask),
        axis=1,
      )
      new_normed_input = revin(new_patched_input, new_mu, new_sigma, reverse=False)
      (_, _, new_normed_output, _), decode_cache = module(
        new_normed_input, new_mask, decode_cache
      )
      new_renormed_output = jax_einshape(
        "bm(oq)->bmoq",
        revin(new_normed_output, new_mu, new_sigma, reverse=True),
        o=module.o,
        q=module.q,
      )[..., -1, :, :]

      return (
        (
          new_renormed_output[..., module.decode_index],
          carry_stats,
          decode_cache,
        ),
        new_renormed_output,
      )

    if num_decode_steps > 0:
      _, ar_renormed_outputs = _ar_decode(
        self,
        (
          renormed_outputs[..., -1, :, self.decode_index],
          (last_n, last_mu, last_sigma),
          decode_cache,
        ),
        jnp.arange(num_decode_steps),
      )
    else:
      ar_renormed_outputs = None

    return renormed_outputs, renormed_quantile_spread, ar_renormed_outputs

  def compile(
    self,
    context: int,
    horizon: int,
    per_core_batch_size: int = 1,
  ):
    if context % self.p != 0:
      logging.info(
        "When compiling, context needs to be multiple of the patch size %d."
        " Modifying context to %d.",
        self.p,
        context := math.ceil(context / self.p) * self.p,
      )
    if horizon % self.o != 0:
      logging.info(
        "When compiling, horizon needs to be multiple of the output patch"
        " size %d. Modifying horizon to %d.",
        self.o,
        horizon := math.ceil(horizon / self.o) * self.o,
      )

    self.context = context
    self.horizon = horizon
    self.per_core_batch_size = per_core_batch_size

    @nnx.pmap(
      in_axes=(None, None, 0, 0),
      out_axes=(0, 0, 0),
      devices=jax.devices(self.backend),
      axis_size=self.num_devices,
      static_broadcasted_argnums=(1,),
      axis_name="global_batch",
    )
    def compiled_decode_kernel(model, horizon, inputs, masks):
      return model.decode(horizon, inputs, masks)

    self.compiled_decode = functools.partial(compiled_decode_kernel, self)


def _flip_quantile_fn(x):
  return jnp.concatenate([x[..., :1], jnp.flip(x[..., 1:], axis=-1)], axis=-1)


@functools.partial(
  jax.jit,
  donate_argnums=(0, 1, 2),
)
def _force_flip_invariance_fn(
  flipped_pf_outputs,
  flipped_quantile_spreads,
  flipped_ar_outputs,
):
  """Forces flip invariance."""
  flipped_pf_outputs = _flip_quantile_fn(flipped_pf_outputs)
  flipped_pf_outputs = jax_einshape("tb...->(tb)...", flipped_pf_outputs)
  flipped_quantile_spreads = _flip_quantile_fn(flipped_quantile_spreads)
  flipped_quantile_spreads = jax_einshape("tb...->(tb)...", flipped_quantile_spreads)
  to_concat = [flipped_pf_outputs[:, -1, ...]]
  if flipped_ar_outputs is not None:
    flipped_ar_outputs = _flip_quantile_fn(flipped_ar_outputs)
    flipped_ar_outputs = jax_einshape("tbno...->(tb)(no)...", flipped_ar_outputs)
    to_concat.append(flipped_ar_outputs)
  flipped_full_forecast = jnp.concatenate(to_concat, axis=1)

  return flipped_quantile_spreads, flipped_pf_outputs, flipped_full_forecast


@functools.partial(
  jax.jit,
  static_argnames=("max_horizon",),
  donate_argnums=(0,),
)
def _use_continuous_quantile_head_fn(full_forecast, quantile_spreads, max_horizon):
  """Uses continuous quantile head."""
  to_stack = [full_forecast[..., :max_horizon, 0]]
  for quantile_index in [1, 2, 3, 4]:
    to_stack.append(
      quantile_spreads[:, :max_horizon, quantile_index]
      - quantile_spreads[:, :max_horizon, 5]
      + full_forecast[:, :max_horizon, 5]
    )
  to_stack.append(full_forecast[..., :max_horizon, 5])
  for quantile_index in [6, 7, 8, 9]:
    to_stack.append(
      quantile_spreads[:, :max_horizon, quantile_index]
      - quantile_spreads[:, :max_horizon, 5]
      + full_forecast[:, :max_horizon, 5]
    )
  return jnp.stack(to_stack, axis=-1)


@functools.partial(jax.jit, donate_argnums=(0,))
def _fix_quantile_crossing_fn(full_forecast):
  """Fixes quantile crossing."""
  lower_quantiles = _scan_along_axis(
    lambda carry, x: (w := jnp.minimum(carry, x), w),
    init=full_forecast[..., 5],
    xs=full_forecast[..., 1:5],
    axis=-1,
    reverse=True,
  )[1]
  upper_quantiles = _scan_along_axis(
    lambda carry, x: (w := jnp.maximum(carry, x), w),
    init=full_forecast[..., 5],
    xs=full_forecast[..., 6:10],
    axis=-1,
    reverse=False,
  )[1]
  return jnp.concatenate(
    [
      full_forecast[..., :1],
      lower_quantiles,
      full_forecast[..., 5:6],
      upper_quantiles,
    ],
    axis=-1,
  )


@functools.partial(jax.jit, static_argnames=("fc",), donate_argnums=(1, 2))
def _before_model_decode(fc, inputs, masks):
  """All Jax steps before model decode call."""
  if fc.infer_is_positive:
    is_positive = jnp.all(inputs >= 0, axis=-1, keepdims=True)
  else:
    is_positive = None

  if fc.normalize_inputs:
    mu = jnp.mean(inputs, axis=-1, keepdims=True)
    sigma = jnp.std(inputs, axis=-1, keepdims=True)
    inputs = revin(inputs, mu, sigma, reverse=False)
  else:
    mu, sigma = None, None

  inputs = jax_einshape("(tb)...->tb...", inputs, b=fc.per_core_batch_size)
  masks = jax_einshape("(tb)...->tb...", masks, b=fc.per_core_batch_size)

  return inputs, masks, is_positive, mu, sigma


@functools.partial(
  jax.jit,
  static_argnames=(
    "fc",
    "p",
  ),
  donate_argnums=(1, 2, 3, 4, 5, 6, 7, 8, 9),
)
def _after_model_decode(
  fc,
  pf_outputs,
  quantile_spreads,
  ar_outputs,
  flipped_pf_outputs,
  flipped_quantile_spreads,
  flipped_ar_outputs,
  is_positive,
  mu,
  sigma,
  p,
):
  """All Jax steps after model decode call."""
  # t: num_devices, b: per_core_batch_size
  pf_outputs = jax_einshape("tb...->(tb)...", pf_outputs)
  quantile_spreads = jax_einshape("tb...->(tb)...", quantile_spreads)
  to_concat = [pf_outputs[:, -1, ...]]
  if ar_outputs is not None:
    ar_outputs = jax_einshape("tbno...->(tb)(no)...", ar_outputs)
    to_concat.append(ar_outputs)
  full_forecast = jnp.concatenate(to_concat, axis=1)

  if fc.force_flip_invariance:
    (
      flipped_quantile_spreads,
      flipped_pf_outputs,
      flipped_full_forecast,
    ) = _force_flip_invariance_fn(
      flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs
    )
    quantile_spreads = (quantile_spreads - flipped_quantile_spreads) / 2
    pf_outputs = (pf_outputs - flipped_pf_outputs) / 2
    full_forecast = (full_forecast - flipped_full_forecast) / 2

  if fc.use_continuous_quantile_head:
    full_forecast = _use_continuous_quantile_head_fn(
      full_forecast, quantile_spreads, fc.max_horizon
    )

  if fc.return_backcast:
    full_backcast = jax_einshape("...npq->...(np)q", pf_outputs[:, :-1, :p, :])
    full_forecast = jnp.concatenate([full_backcast, full_forecast], axis=1)

  if fc.fix_quantile_crossing:
    full_forecast = _fix_quantile_crossing_fn(full_forecast)

  if fc.normalize_inputs:
    full_forecast = revin(full_forecast, mu, sigma, reverse=True)

  if is_positive is not None:
    full_forecast = jnp.where(
      is_positive[..., None],
      jnp.maximum(full_forecast, jnp.zeros_like(full_forecast)),
      full_forecast,
    )

  return full_forecast


class TimesFM_2p5_200M_flax(timesfm_2p5_base.TimesFM_2p5):
  """Flax implementation of TimesFM 2.5 with 200M parameters."""

  model: nnx.Module = TimesFM_2p5_200M_flax_module()

  @classmethod
  def from_pretrained(
    cls,
    model_id: str = "google/timesfm-2.5-200m-flax",
    *,
    revision: str | None = None,
    cache_dir: str | Path | None = None,
    force_download: bool = False,
    proxies: Dict | None = None,
    resume_download: bool | None = None,
    local_files_only: bool | None = None,
    token: str | None = None,
    **model_kwargs,
  ):
    """Loads a Flax TimesFM model."""

    # Create an instance of the model wrapper class.
    instance = cls(**model_kwargs)

    # Determine the path to the model weights.
    model_file_path = ""
    if os.path.isdir(model_id):
      logging.info("Loading checkpoint from local directory: %s", model_id)
      model_file_path = model_id
    else:
      logging.info("Downloading checkpoint from Hugging Face repo %s", model_id)
      model_file_path = huggingface_hub.snapshot_download(
        repo_id=model_id,
        revision=revision,
        cache_dir=cache_dir,
        force_download=force_download,
        proxies=proxies,
        resume_download=resume_download,
        token=token,
        local_files_only=local_files_only,
      )
      logging.info("Loading checkpoint from: %s", model_file_path)

    checkpointer = ocp.StandardCheckpointer()
    graph, state = nnx.split(instance.model)
    state = checkpointer.restore(model_file_path, state)
    instance.model = nnx.merge(graph, state)
    return instance

  def compile(
    self,
    forecast_config: configs.ForecastConfig,
    dryrun: bool = True,
    **kwargs
  ):
    # Acrobym used during validation.
    print("Compiling model...")

    fc = forecast_config
    if fc.max_context % self.model.p != 0:
      logging.info(
        "When compiling, max context needs to be multiple of the patch size"
        " %d. Using max context = %d instead.",
        self.model.p,
        new_context := math.ceil(fc.max_context / self.model.p) * self.model.p,
      )
      fc = dataclasses.replace(fc, max_context=new_context)
    if fc.max_horizon % self.model.o != 0:
      logging.info(
        "When compiling, max horizon needs to be multiple of the output patch"
        " size %d. Using max horizon = %d instead.",
        self.model.o,
        new_horizon := math.ceil(fc.max_horizon / self.model.o) * self.model.o,
      )
      fc = dataclasses.replace(fc, max_horizon=new_horizon)
    if fc.max_context + fc.max_horizon > self.model.config.context_limit:
      raise ValueError(
        "Context + horizon must be less than the context limit."
        f" {fc.max_context} + {fc.max_horizon} >"
        f" {self.model.config.context_limit}."
      )
    if fc.use_continuous_quantile_head and (fc.max_horizon > self.model.os):
      raise ValueError(
        f"Continuous quantile head is not supported for horizons > {self.model.os}."
      )

    self.forecast_config = fc
    self.model.compile(
      context=self.forecast_config.max_context,
      horizon=self.forecast_config.max_horizon,
      per_core_batch_size=fc.per_core_batch_size,
    )
    self.per_core_batch_size = self.forecast_config.per_core_batch_size
    self.num_devices = self.model.num_devices
    self.global_batch_size = (
      self.forecast_config.per_core_batch_size * self.model.num_devices
    )

    def compiled_decode_kernel(fc, horizon, inputs, masks):
      inputs = jnp.array(inputs, dtype=jnp.float32)
      masks = jnp.array(masks, dtype=jnp.bool)
      if horizon > fc.max_horizon:
        raise ValueError(
          f"Horizon must be less than the max horizon. {horizon} > {fc.max_horizon}."
        )
      to_trim = fc.max_horizon - horizon

      inputs, masks, is_positive, mu, sigma = _before_model_decode(fc, inputs, masks)

      pf_outputs, quantile_spreads, ar_outputs = self.model.compiled_decode(
        fc.max_horizon, inputs, masks
      )
      if fc.force_flip_invariance:
        flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = (
          self.model.compiled_decode(fc.max_horizon, -inputs, masks)
        )
      else:
        flipped_pf_outputs, flipped_quantile_spreads, flipped_ar_outputs = (
          None,
          None,
          None,
        )

      full_forecast = _after_model_decode(
        fc,
        pf_outputs,
        quantile_spreads,
        ar_outputs,
        flipped_pf_outputs,
        flipped_quantile_spreads,
        flipped_ar_outputs,
        is_positive,
        mu,
        sigma,
        self.model.p,
      )
      full_forecast_np = np.array(full_forecast)
      del full_forecast
      try_gc()
      if to_trim > 0:
        full_forecast_np = full_forecast_np[..., :-to_trim, :]
      return full_forecast_np[..., 5], full_forecast_np

    self.compiled_decode = functools.partial(
      compiled_decode_kernel, self.forecast_config
    )

    if dryrun:
      _ = self.compiled_decode(
        self.forecast_config.max_horizon,
        jnp.zeros(
          (self.global_batch_size, self.forecast_config.max_context), dtype=jnp.float32
        ),
        jnp.zeros(
          (self.global_batch_size, self.forecast_config.max_context), dtype=jnp.bool
        ),
      )
    print("Compiling done.")
