import spaces                        # for ZeroGPU support
import gradio as gr
import pandas as pd
import numpy as np
import torch
import subprocess
from threading import Thread
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    AutoProcessor,
    TextIteratorStreamer
)

# ─── MODEL SETUP ────────────────────────────────────────────────────────────────
MODEL_NAME = "bytedance-research/ChatTS-14B"

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME, trust_remote_code=True
)
processor = AutoProcessor.from_pretrained(
    MODEL_NAME, trust_remote_code=True, tokenizer=tokenizer
)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    device_map="auto",
    torch_dtype=torch.float16
)
model.eval()

# ─── HELPER FUNCTIONS ──────────────────────────────────────────────────────────

def create_default_timeseries():
    """Create default time series with sudden increase"""
    seq_len = 256
    y = np.zeros(seq_len, dtype=np.float32)
    y[100:] += 1
    df = pd.DataFrame({"default_series": y})
    return df

def process_csv_file(csv_file):
    """Process CSV file and return DataFrame with validation"""
    if csv_file is None:
        return None, "No file uploaded"
    
    try:
        df = pd.read_csv(csv_file.name)
        
        # drop columns with empty names or all-NaNs
        df.columns = [str(c).strip() for c in df.columns]
        df = df.loc[:, [c for c in df.columns if c]]
        df = df.dropna(axis=1, how="all")
        print(f"File {csv_file.name} loaded. {df.columns=}")

        if df.shape[1] == 0:
            return None, "No valid time-series columns found."
        if df.shape[1] > 15:
            return None, f"Too many series ({df.shape[1]}). Max allowed = 15."

        # Validate ALL columns as time series
        ts_names, ts_list = [], []
        for name in df.columns:
            series = df[name]
            # ensure float dtype
            if not pd.api.types.is_float_dtype(series):
                try:
                    series = pd.to_numeric(series, errors='coerce')
                except:
                    return None, f"Series '{name}' cannot be converted to float type."
            
            # trim trailing NaNs only
            last_valid = series.last_valid_index()
            if last_valid is None:
                continue
            trimmed = series.loc[:last_valid].to_numpy(dtype=np.float32)
            length = trimmed.shape[0]
            if length < 64 or length > 1024:
                return None, f"Series '{name}' length {length} invalid. Must be 64 to 1024."
            ts_names.append(name)
            ts_list.append(trimmed)

        if not ts_list:
            return None, "All time series are empty after trimming NaNs."
        print(f"Successfully loaded {len(ts_names)} time series: {', '.join(ts_names)}")
            
        return df, f"Successfully loaded {len(ts_names)} time series: {', '.join(ts_names)}"
        
    except Exception as e:
        return None, f"Error processing file: {str(e)}"

def preview_csv(csv_file, use_default):
    """Preview uploaded CSV file immediately"""
    if csv_file is None:
        return gr.LinePlot(value=pd.DataFrame()), "Please upload a CSV file first", gr.Dropdown(), False
    
    df, message = process_csv_file(csv_file)
    
    if df is None:
        return gr.LinePlot(value=pd.DataFrame()), message, gr.Dropdown(), False
    
    # Create dropdown choices
    column_choices = list(df.columns)
    
    # Create plot with first column as default
    first_column = column_choices[0]
    df_with_index = df.copy()
    df_with_index["_internal_idx"] = np.arange(len(df[first_column].values))
    plot = gr.LinePlot(
        df_with_index,
        x="_internal_idx",
        y=first_column,
        title=f"Time Series: {first_column}"
    )
    
    # Update dropdown
    dropdown = gr.Dropdown(
        choices=column_choices,
        value=first_column,
        label="Select a Column to Visualize"
    )
    
    print("Successfully generated preview!")
    
    return plot, message, dropdown, False  # Set use_default to False when file is uploaded

def clear_csv():
    """Clear uploaded CSV file immediately"""
    df, message = process_csv_file(None)
    
    return gr.LinePlot(value=pd.DataFrame()), message, gr.Dropdown()


def update_plot(csv_file, selected_column):
    """Update plot based on selected column"""
    if csv_file is None or selected_column is None:
        return gr.LinePlot(value=pd.DataFrame())
    
    df, _ = process_csv_file(csv_file)
    if df is None:
        return gr.LinePlot(value=pd.DataFrame())
    
    df_with_index = df.copy()
    df_with_index["_internal_idx"] = np.arange(len(df[selected_column].values))
    
    plot = gr.LinePlot(
        df_with_index,
        x="_internal_idx",
        y=selected_column,
        title=f"Time Series: {selected_column}"
    )
    
    return plot

def initialize_interface():
    """Initialize interface with default time series"""
    df = create_default_timeseries()
    column_choices = list(df.columns)
    first_column = column_choices[0]
    
    df_with_index = df.copy()
    df_with_index["_internal_idx"] = np.arange(len(df[first_column].values))
    
    plot = gr.LinePlot(
        df_with_index,
        x="_internal_idx",
        y=first_column,
        title=f"Time Series: {first_column}"
    )
    
    dropdown = gr.Dropdown(
        choices=column_choices,
        value=first_column,
        label="Select a Column to Visualize"
    )
    
    message = "Using default time series with sudden increase at step 100"
    
    return plot, message, dropdown, True  # Set use_default to True on initialization

# ─── INFERENCE + VALIDATION ────────────────────────────────────────────────────

@spaces.GPU  # dynamically allocate & release a ZeroGPU device on each call
def infer_chatts_stream(prompt: str, csv_file, use_default):
    """
    Streaming version of ChatTS inference
    """
    print("Start inferring!!!")
    
    if not prompt.strip():
        yield "Please enter a prompt"
        return
    
    # Use default if no file uploaded and use_default is True
    if csv_file is None and use_default:
        df = create_default_timeseries()
        error_msg = None
    else:
        df, error_msg = process_csv_file(csv_file)
    
    if df is None:
        yield "Please upload a CSV file first or the file contains errors"
        return
    
    try:
        # Prepare time series data - use ALL columns
        ts_names, ts_list = [], []
        for name in df.columns:
            series = df[name]
            last_valid = series.last_valid_index()
            if last_valid is not None:
                trimmed = series.loc[:last_valid].to_numpy(dtype=np.float32)
                ts_names.append(name)
                ts_list.append(trimmed)
        
        if not ts_list:
            yield "No valid time series data found. Please upload time series first."
            return
        
        # Clean prompt
        clean_prompt = prompt.replace("<ts>", "").replace("<ts/>", "")

        # Build prompt prefix
        prefix = f"I have {len(ts_list)} time series:\n"
        for name, arr in zip(ts_names, ts_list):
            prefix += f"The {name} is of length {len(arr)}: <ts><ts/>\n"
        
        full_prompt = f"<|im_start|>system\nYou are a helpful assistant. Your name is ChatTS. You can analyze time series data and provide insights. If user asks who you are, you should give your name and capabilities in the language of the prompt. If no time series are provided, you should say 'I cannot answer this question as you haven't provide the timeseries I need' in the language of the prompt. Always check if the user has provided at least one time series data before answering.<|im_end|><|im_start|>user\n{prefix}{clean_prompt}<|im_end|><|im_start|>assistant\n"

        print(f"[debug] {full_prompt}. {len(ts_list)=}, {[len(item) for item in ts_list]=}")

        # Encode inputs
        inputs = processor(
            text=[full_prompt],
            timeseries=ts_list,
            padding=True,
            return_tensors="pt"
        )
        inputs = {k: v.to(model.device) for k, v in inputs.items()}

        if inputs['timeseries'] is not None:
            print(f"[debug] {inputs['timeseries'].shape=}")

        # Generate with streaming
        streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
        inputs.update({
            "max_new_tokens": 512,
            "streamer": streamer,
            "temperature": 0.3
        })
        thread = Thread(
            target=model.generate,
            kwargs=inputs
        )
        thread.start()

        model_output = ""
        for new_text in streamer:
            model_output += new_text
            yield model_output
        
    except Exception as e:
        yield f"Error during inference: {str(e)}"

# ─── GRADIO APP ────────────────────────────────────────────────────────────────

with gr.Blocks(title="ChatTS Demo") as demo:
    gr.Markdown("## ChatTS: Time Series Understanding and Reasoning")
    gr.HTML("""<div style="display:flex;justify-content: center">
<a href="https://github.com/NetmanAIOps/ChatTS"><img alt="github" src="https://img.shields.io/badge/Code-GitHub-blue"></a>
<a href="https://huggingface.co/bytedance-research/ChatTS-14B"><img alt="github" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Model-FFD21E"></a>
<a href="https://arxiv.org/abs/2412.03104"><img alt="preprint" src="https://img.shields.io/static/v1?label=arXiv&amp;message=2412.03104&amp;color=B31B1B&amp;logo=arXiv"></a>
</div>""")
    gr.Markdown("Try ChatTS with the default time series, or upload a CSV file (Example: [ts_example.csv](https://github.com/NetManAIOps/ChatTS/blob/main/demo/ts_example.csv)) containing UTS/MTS where each column is a dimension (no index column). All columns will be used as input of ChatTS automatically.")

    # State to track whether to use default time series
    use_default_state = gr.State(value=True)

    with gr.Row():
        with gr.Column(scale=1):
            upload = gr.File(
                label="Upload CSV File",
                file_types=[".csv"],
                type="filepath"
            )
            
            prompt_input = gr.Textbox(
                lines=6,
                placeholder="Enter your question here...",
                label="Analysis Prompt",
                value="Please analyze all the given time series and provide insights about the local fluctuations in the time series in detail."
            )
            
            run_btn = gr.Button("Run ChatTS", variant="primary")
            
        with gr.Column(scale=2):
            series_selector = gr.Dropdown(
                label="Select a Column to Visualize",
                choices=[],
                value=None
            )
            plot_out = gr.LinePlot(value=pd.DataFrame(), label="Time Series Visualization")
            file_status = gr.Textbox(
                label="File Status", 
                interactive=False,
                lines=2
            )

    text_out = gr.Textbox(
        lines=10, 
        label="ChatTS Analysis Results",
        interactive=False
    )

    # Initialize interface with default data
    demo.load(
        fn=initialize_interface,
        outputs=[plot_out, file_status, series_selector, use_default_state]
    )

    # Event handlers
    upload.upload(
        fn=preview_csv,
        inputs=[upload, use_default_state],
        outputs=[plot_out, file_status, series_selector, use_default_state]
    )

    upload.clear(
        fn=clear_csv,
        inputs=[],
        outputs=[plot_out, file_status, series_selector]
    )

    series_selector.change(
        fn=update_plot,
        inputs=[upload, series_selector],
        outputs=[plot_out]
    )
    
    run_btn.click(
        fn=infer_chatts_stream,
        inputs=[prompt_input, upload, use_default_state],
        outputs=[text_out]
    )

if __name__ == '__main__':
    demo.launch()