import torch
import torch.nn as nn


class BaseRNN(nn.Module):
    def __init__(self, time_in, feat_in, context_in, hidden_size,
                 rnn_size, out_size, rnn_type="RNN", device="cpu"):
        super().__init__()

        self.time_proj = nn.Linear(time_in, hidden_size, device=device)
        self.feat_proj = nn.Linear(feat_in, hidden_size, device=device)
        self.context_proj = (
            nn.Linear(context_in, hidden_size, device=device) if context_in is not None else None
        )

        if rnn_type == "RNN":
            self.rnn = nn.RNN(
                input_size=hidden_size*3,
                hidden_size=rnn_size,
                num_layers=2,
                batch_first=True,
                device=device
            )
        elif rnn_type == "LSTM":
            self.rnn = nn.LSTM(
                input_size=hidden_size*3,
                hidden_size=rnn_size,
                num_layers=2,
                batch_first=True,
                device=device
            )
        else:
            raise ValueError("Unsupported rnn_type")

        readout_in = rnn_size
        self.readout = nn.Linear(readout_in, out_size, device=device)

        self.context_in = context_in
        self.rnn_size = rnn_size

    def forward(self, time_seq, feat_init=None, context=None,
                init_steps=0, pred_steps=0, offset=0, hidden=None):

        b, total_len, _ = time_seq.shape
        assert init_steps + offset <= total_len

        # Project everything once
        time_proj = self.time_proj(time_seq)                    # (b, T, hidden)
        if context is not None:
            assert self.context_proj is not None
            ctx = self.context_proj(context).unsqueeze(1)       # (b,1,hidden)
            ctx = ctx.expand(-1, total_len, -1)                 # broadcast
        else:
            ctx = torch.zeros(b, total_len, time_proj.size(-1),
                               device=time_seq.device)

        feat_proj_or_zeros = torch.zeros(b, total_len, time_proj.size(-1), device=time_seq.device)
        if feat_init is not None:
            feat_proj = self.feat_proj(feat_init)               # (b, init_steps, hidden)
            # Insert feat input *only* in first init_steps
            feat_proj_or_zeros[:, :init_steps, :] += feat_proj
        
        rnn_input = torch.cat([time_proj, ctx, feat_proj_or_zeros], dim=-1)

        # Run the RNN over full sequence
        outputs, hidden_next = self.rnn(rnn_input, hidden)

        # Prediction window
        pred_h = outputs[:, init_steps:init_steps + pred_steps]   # (b, pred, rnn_size)

        # Build readout input
        time_raw = time_seq[:, init_steps:init_steps + pred_steps]

        preds = self.readout(pred_h)
        return preds, hidden_next



            

class TwoInputRNN(BaseRNN):
    """
    Many-to-many RNN with an initialization phase that consumes both:
      - time inputs (always present)
      - feature inputs (only during init_steps)
    After init_steps, only time inputs are provided and the model predicts a sequence.
    Predictions can be compared to targets with a specified `offset`.
    """
    def __init__(self, time_in, feat_in, hidden_size, rnn_size, out_size, device="cpu"):
        super().__init__(time_in, feat_in, context_in=None, hidden_size=hidden_size, rnn_size=rnn_size, out_size=out_size, device=device)

class ThreeInputRNN(BaseRNN):
    """
    Many-to-many RNN with an initialization phase that consumes three inputs:
      - time inputs (always present)
      - feature inputs (only during init_steps)
      - context inputs (always present + constant over time)
    After init_steps, only time and context inputs are provided and the model predicts a sequence.
    Predictions can be compared to targets with a specified `offset`.
    """
    def __init__(self, time_in, feat_in, context_in, hidden_size, rnn_size, out_size, device="cpu"):
        super().__init__(time_in, feat_in, context_in=context_in, hidden_size=hidden_size, rnn_size=rnn_size, out_size=out_size, device=device)

class TwoInputLSTM(BaseRNN):
    """
    Many-to-many LSTM with an initialization phase that consumes both:
      - time inputs (always present)
      - feature inputs (only during init_steps)
    After init_steps, only time inputs are provided and the model predicts a sequence.
    Predictions can be compared to targets with a specified `offset`.
    """
    def __init__(self, time_in, feat_in, hidden_size, rnn_size, out_size, device="cpu"):
        super().__init__(time_in, feat_in, context_in=None, hidden_size=hidden_size, rnn_size=rnn_size, out_size=out_size, rnn_type="LSTM", device=device)

class ThreeInputLSTM(BaseRNN):
    """
    Many-to-many LSTM with an initialization phase that consumes three inputs:
      - time inputs (always present)
      - feature inputs (only during init_steps)
      - context inputs (always present + constant over time)
    After init_steps, only time and context inputs are provided and the model predicts a sequence.
    Predictions can be compared to targets with a specified `offset`.
    """
    def __init__(self, time_in, feat_in, context_in, hidden_size, rnn_size, out_size, device="cpu"):
        super().__init__(time_in, feat_in, context_in=context_in, hidden_size=hidden_size, rnn_size=rnn_size, out_size=out_size, rnn_type="LSTM", device=device)