import time
import asyncio
import logging
import functools
from typing import Coroutine
import numpy as np
import pandas as pd

from nirvana.executors.llm_backbone import LLMClient
from nirvana.lineage.abstractions import LineageNode, OpOutputsType

logger = logging.getLogger(__name__)


ORDERED_MODELS= [  # llm list ordered by speed and intelligence (weak but fast -> strong but slow)
    "gpt-4.1-nano-2025-04-14",
    "gpt-4o-mini-2024-07-18",
    "gpt-4.1-mini-2025-04-14",
    "gpt-4.1-2025-04-14",
]


def concate_output(output1: OpOutputsType, output2: OpOutputsType):
    return output1 + output2


def cosine_similarity(x: np.array, y: np.array):
    x_norms = np.linalg.norm(x, axis=1)
    y_norms = np.linalg.norm(y, axis=1)
    dot_products = np.sum(x * y, axis=1)
    return dot_products / (x_norms * y_norms)


class PhysicalOptimizer:
    def __init__(
        self,
        agent: LLMClient = None,
        available_models: list[str] = ORDERED_MODELS,
        num_samples: int = 10,
        improve_margin: float = 0.2,
    ):
        self.agent = agent
        self.available_models = available_models if available_models else ORDERED_MODELS
        self.num_samples = num_samples
        self.improve_margin = improve_margin

    def should_optimize(self, op_name: str, input_data: pd.DataFrame | list[pd.DataFrame], num_samples: int) -> bool:
        if op_name == "join":
            # and len(input_data[0]) * len(input_data[1]) < 2 * num_samples:
            # we don't optimize join yet
            return False
        elif op_name in ["filter", "map"] and len(input_data) < 2 * num_samples:
            return False
        elif op_name in ["reduce", "rank"]:
            return False
        else:
            return True

    def split_input_data(self, op_name: str, input_data: pd.DataFrame | list[pd.DataFrame], num_samples: int):
        if op_name == "join":
            left_length, right_length = len(input_data[0]), len(input_data[1])
            if left_length > right_length:
                return [input_data[0].iloc[:num_samples], input_data[1]], [input_data[0].iloc[num_samples:], input_data[1]]
            else:
                return [input_data[0], input_data[1].iloc[:num_samples]], [input_data[0], input_data[1].iloc[num_samples:]]
        else:
            return input_data.iloc[:num_samples], input_data.iloc[num_samples:]
        
    def _serialize_map_outputs(self, outputs: dict[str, list]) -> list[str]:
        serialized_outputs = []
        for values in zip(*outputs.values()):
            values_with_key = map(lambda k, v: f"{k}: {v}", outputs.keys(), values)
            serialized_outputs.append(", ".join(values_with_key))
        return serialized_outputs

    async def optimize_exec_model(
        self, 
        node: LineageNode, 
        input_data: pd.DataFrame | list[pd.DataFrame],
    ):
        if not self.should_optimize(node.op_name, input_data, self.num_samples):
            node.operator.model = self.available_models[0] if node.op_name in ["join", "rank"] else self.available_models[-1]
            node_output = await node.execute_operation(input_data)
            return node_output

        optimize_start_time = time.time()
        train_set, test_set = self.split_input_data(node.op_name, input_data, self.num_samples)
        exec_model = self.available_models[0]
        node.operator.model = exec_model 
        node_output = None

        if node.op_name == "filter":
            worst_output = await node.execute_operation(train_set)
            node_output = worst_output
            
            improve_score_list = [0.0]
            mismatch_idx_worst2subworst = None
            mismatch_idx_subworst2subbest = None
            for model_id in range(1, len(ORDERED_MODELS)):
                if model_id == 1:
                    better_model = ORDERED_MODELS[model_id]
                    node.operator.model = better_model
                    better_output = await node.execute_operation(train_set)
                    node_output.cost += better_output.cost
                    subworst_output = np.array(better_output.output)
                    mismatch_idx_worst2subworst = np.array(worst_output.output) != subworst_output
                    mismatch_record_ratio = sum(mismatch_idx_worst2subworst) / len(worst_output.output)
                    improve_score_list.append(mismatch_record_ratio)

                elif model_id == 2:
                    better_model = ORDERED_MODELS[model_id]
                    node.operator.model = better_model
                    better_output = await node.execute_operation(train_set)
                    node_output.cost += better_output.cost
                    subbest_output = np.array(better_output.output)
                    mismatch_idx_subworst2subbest = subworst_output != subbest_output
                    if mismatch_idx_worst2subworst.all():
                        mismatch_ratio = 0.0
                    else:
                        mismatch_ratio = (
                            (mismatch_idx_subworst2subbest & ~mismatch_idx_worst2subworst).sum() / (~mismatch_idx_worst2subworst).sum()
                        )
                    improve_score = improve_score_list[-1] + mismatch_ratio * (1 - improve_score_list[-1])
                    improve_score_list.append(improve_score)

                elif model_id == 3:
                    better_model = ORDERED_MODELS[model_id]
                    node.operator.model = better_model
                    all_matched_idx = ~mismatch_idx_subworst2subbest & ~mismatch_idx_worst2subworst
                    if all_matched_idx.any():
                        better_output = await node.execute_operation(train_set.iloc[all_matched_idx])
                        node_output.cost += better_output.cost
                        best_output = np.array(better_output.output)
                        mismatch_idx_subbest2best = subbest_output[all_matched_idx] != best_output
                        mismatch_ratio = mismatch_idx_subbest2best.sum() / all_matched_idx.sum()
                        subbest_output[all_matched_idx] = best_output
                        better_output.output = subbest_output.tolist()
                    else:
                        mismatch_ratio = 0.0
                    if mismatch_idx_worst2subworst.any():
                        matched_ratio = (
                            (~mismatch_idx_subworst2subbest & mismatch_idx_worst2subworst).sum() / mismatch_idx_worst2subworst.sum()
                        )
                    else:
                        matched_ratio = 0.0
                    improve_score = mismatch_ratio * (1 - improve_score_list[-1]) + improve_score_list[-1] - improve_score_list[-2] + matched_ratio
                    improve_score_list.append(improve_score)

                if improve_score_list[-1] - improve_score_list[-2] > self.improve_margin:
                    exec_model = better_model
                    node_output.output = better_output.output

        elif node.op_name == "map":
            worst_output = await node.execute_operation(train_set)
            node_output = worst_output
            serialized_outputs = self._serialize_map_outputs(worst_output.outputs)
            worst_output_embeds, embed_cost = await self.agent.create_embedding(serialized_outputs)
            node_output.cost += embed_cost

            improve_score_list = [0.0]
            mismatch_idx_worst2subworst = None
            mismatch_idx_subworst2subbest = None
            for model_id in range(1, len(ORDERED_MODELS)):
                if model_id == 1:
                    better_model = ORDERED_MODELS[model_id]
                    node.operator.model = better_model
                    better_output = await node.execute_operation(train_set)
                    node_output.cost += better_output.cost
                    serialized_outputs = self._serialize_map_outputs(better_output.outputs)
                    subworst_output_embeds, embed_cost = await self.agent.create_embedding(serialized_outputs)
                    node_output.cost += embed_cost
                    cos_sim = cosine_similarity(worst_output_embeds, subworst_output_embeds)
                    mismatch_idx_worst2subworst = cos_sim < 0.5
                    mismatch_ratio = sum(mismatch_idx_worst2subworst) / len(serialized_outputs)
                    improve_score_list.append(mismatch_ratio)
                
                elif model_id == 2:
                    better_model = ORDERED_MODELS[model_id]
                    node.operator.model = better_model
                    better_output = await node.execute_operation(train_set)
                    node_output.cost += better_output.cost
                    subbest_output = np.array(list(better_output.outputs.values()))
                    serialized_outputs = self._serialize_map_outputs(better_output.outputs)
                    subbest_output_embeds, embed_cost = await self.agent.create_embedding(serialized_outputs)
                    node_output.cost += embed_cost
                    cos_sim = cosine_similarity(subworst_output_embeds, subbest_output_embeds)
                    mismatch_idx_subworst2subbest = cos_sim < 0.5
                    if mismatch_idx_worst2subworst.all():
                        mismatch_ratio = 0.0
                    else:
                        mismatch_ratio = (
                            (mismatch_idx_subworst2subbest & ~mismatch_idx_worst2subworst).sum() / (~mismatch_idx_worst2subworst).sum()
                        )
                    improve_score = improve_score_list[-1] + mismatch_ratio * (1 - improve_score_list[-1])
                    improve_score_list.append(improve_score)
                
                elif model_id == 3:
                    better_model = ORDERED_MODELS[model_id]
                    node.operator.model = better_model
                    all_matched_idx = ~mismatch_idx_subworst2subbest & ~mismatch_idx_worst2subworst
                    if all_matched_idx.any():
                        better_output = await node.execute_operation(train_set.iloc[all_matched_idx])
                        node_output.cost += better_output.cost
                        best_output = np.array(list(better_output.outputs.values()))
                        serialized_outputs = self._serialize_map_outputs(better_output.outputs)
                        best_output_embeds, embed_cost = await self.agent.create_embedding(serialized_outputs)
                        node_output.cost += embed_cost
                        cos_sim = cosine_similarity(subbest_output_embeds[all_matched_idx], best_output_embeds)
                        mismatch_idx_subbest2best = cos_sim < 0.5
                        mismatch_ratio = mismatch_idx_subbest2best.sum() / all_matched_idx.sum()
                        subbest_output[:, all_matched_idx] = best_output
                        for idx, key in enumerate(better_output.outputs.keys()):
                            better_output.outputs[key] = subbest_output[idx].tolist()
                    else:
                        mismatch_ratio = 0.0
                    if mismatch_idx_worst2subworst.any():
                        matched_ratio = (
                            (~mismatch_idx_subworst2subbest & mismatch_idx_worst2subworst).sum() / mismatch_idx_worst2subworst.sum()
                        )
                    else:
                        matched_ratio = 0.0
                    improve_score = mismatch_ratio * (1 - improve_score_list[-1]) + improve_score_list[-1] - improve_score_list[-2] + matched_ratio
                    improve_score_list.append(improve_score)

                if improve_score_list[-1] - improve_score_list[-2] > self.improve_margin:
                    exec_model = better_model
                    node_output.outputs = better_output.outputs
        
        optimize_end_time = time.time()
        logger.info(f"Physical Plan Optimization Time: {optimize_end_time - optimize_start_time:.4f} sec")

        node.operator.model = exec_model
        rest_output = await node.execute_operation(test_set)
        node_output = concate_output(node_output, rest_output)
        return node_output

    async def _optimize(self, node: LineageNode) -> tuple[pd.DataFrame, float]:
        if node.left_child:
            dataframe_from_left_node, cost_from_left_branch = await self._optimize(node.left_child)
        if node.right_child:
            dataframe_from_right_node, cost_from_right_branch = await self._optimize(node.right_child)

        if node.op_name == "scan":
            output_from_node = await node.run()
            return output_from_node.output, output_from_node.cost
        
        elif node.op_name == "join":
            token_cost = cost_from_left_branch + cost_from_right_branch
            output_from_node = await self.optimize_exec_model(node, [dataframe_from_left_node, dataframe_from_right_node])
            dataframe_from_node = await node.collate_dataframe([dataframe_from_left_node, dataframe_from_right_node], output_from_node)
            token_cost += output_from_node.cost
            return dataframe_from_node, token_cost
        
        else:
            token_cost = cost_from_left_branch
            output_from_node = await self.optimize_exec_model(node, dataframe_from_left_node)
            dataframe_from_node = await node.collate_dataframe(dataframe_from_left_node, output_from_node)
            token_cost += output_from_node.cost
            return dataframe_from_node, token_cost

    def optimize(
        self, 
        plan: LineageNode,
    ):
        execution_start_time = time.time()
        dataframe_from_node, total_token_cost = asyncio.run(self._optimize(plan))
        execution_end_time = time.time()
        execution_time = execution_end_time - execution_start_time
        return dataframe_from_node, total_token_cost, execution_time
