import dataiku
import pandas as pd
from copy import deepcopy

from dku_utils.projects.project_commons import get_current_project_and_variables
from parameters_control_utils import insight_builder, is_under_control, make_metric_block, generate_capability_metrics_data
import insights_static
import logging

logger = logging.getLogger(__name__)

client = dataiku.api_client()
project, variables = get_current_project_and_variables()
selected_variables = variables["standard"]["selected_parameterz"]
using_pa = variables["standard"].get("using_pa", "")
chart_type = variables["standard"].get("chart_type", "X_Bar")
subgroup_size = variables["standard"].get("subgroup_size", "")


input_dku_dataset = dataiku.Dataset("selected_study_data_enhanced")
df_raw = input_dku_dataset.get_dataframe()
limits_row = input_dku_dataset.get_dataframe(limit=1)

# Group by subgroup
agg_funcs = {col: ["mean", lambda x: x.max() - x.min()] for col in df_raw.select_dtypes(include="number").columns if col != "_subgroup"}

df_by_subgroup = df_raw.groupby("_subgroup").mean(numeric_only=True).reset_index()

# Generate Capability Metrics
capability_metrics_df = generate_capability_metrics_data(df_raw, selected_variables, chart_type)
capability_metrics_output = dataiku.Dataset("capability_metrics")
capability_metrics_output.write_with_schema(capability_metrics_df)

available_capability_metrics = capability_metrics_df.columns.tolist()

# Delete all the existing insights:
current_insights_ids = [item["id"] for item in project.list_insights() if "id" in item]
for id in current_insights_ids:
    insight = project.get_insight(id)
    insight.delete()

insights_ids = {}


for variable_name in selected_variables:

    # Get creation info from insights_static dynamically
    x_bar_info = getattr(insights_static, "x_bar_creation_info")
    r_chart_info = getattr(insights_static, "r_chart_creation_info")
    dist_plot_creation_info = getattr(insights_static, "dist_plot_creation_info")
    metrics_creation_info = getattr(insights_static, "metrics_creation_info")
    generic_measure_template = getattr(insights_static, "generic_measure_template")
    target_line = getattr(insights_static, "target_line")
    dist_plot_generic_measures_template = getattr(insights_static, "dist_plot_generic_measures_template")

    # Create appropriate control charts based on chart_type
    if chart_type == "X_Bar":
        # X̅–R Chart
        insights_ids[f"x_bar_{variable_name}_id"] = insight_builder(
            project,
            x_bar_info,
            variable_name,
            "X-bar",
            [f"UCL_Xbar_{variable_name}", f"LCL_Xbar_{variable_name}", f"mean_{variable_name}"],
            reference_column=f"mean_{variable_name}",
            custom_colors={f"m_UCL_Xbar_{variable_name}_AVG_1": "#d62728", f"m_LCL_Xbar_{variable_name}_AVG_1": "#d62728", f"m_mean_{variable_name}_AVG_1": "#1f77b4"},
        )

        insights_ids[f"variability_{variable_name}_id"] = insight_builder(
            project,
            r_chart_info,
            variable_name,
            "R",
            [f"UCL_R_{variable_name}", f"LCL_R_{variable_name}", f"range_{variable_name}"],
            reference_column=f"range_{variable_name}",
            custom_colors={f"m_UCL_R_{variable_name}_AVG_1": "#d62728", f"m_LCL_R_{variable_name}_AVG_1": "#d62728", f"m_range_{variable_name}_AVG_1": "#ff7f0e"},
        )

    elif chart_type == "XS_Bar":
        # X̅–S Chart
        insights_ids[f"x_bar_{variable_name}_id"] = insight_builder(
            project,
            x_bar_info,
            variable_name,
            "X-bar",
            [f"UCL_Xbar_{variable_name}", f"LCL_Xbar_{variable_name}", f"mean_{variable_name}"],
            reference_column=f"mean_{variable_name}",
            custom_colors={f"m_UCL_Xbar_{variable_name}_AVG_1": "#d62728", f"m_LCL_Xbar_{variable_name}_AVG_1": "#d62728", f"m_mean_{variable_name}_AVG_1": "#1f77b4"},
        )

        insights_ids[f"variability_{variable_name}_id"] = insight_builder(
            project,
            r_chart_info,  # Reusing r_chart_info structure for S chart
            variable_name,
            "S",
            [f"UCL_S_{variable_name}", f"LCL_S_{variable_name}", f"std_{variable_name}"],
            reference_column=f"std_{variable_name}",
            custom_colors={f"m_UCL_S_{variable_name}_AVG_1": "#d62728", f"m_LCL_S_{variable_name}_AVG_1": "#d62728", f"m_std_{variable_name}_AVG_1": "#ff7f0e"},
        )

    elif chart_type == "I_MR":
        # I-MR Chart
        insights_ids[f"x_bar_{variable_name}_id"] = insight_builder(
            project,
            x_bar_info,
            variable_name,
            "I (Individual)",
            [f"UCL_I_{variable_name}", f"LCL_I_{variable_name}", variable_name],  # Individual values
            reference_column=variable_name,
            custom_colors={f"m_UCL_I_{variable_name}_AVG_1": "#d62728", f"m_LCL_I_{variable_name}_AVG_1": "#d62728", f"m_{variable_name}_AVG_1": "#1f77b4"},
        )

        insights_ids[f"variability_{variable_name}_id"] = insight_builder(
            project,
            r_chart_info,  # Reusing r_chart_info structure for MR chart
            variable_name,
            "MR (Moving Range)",
            [f"UCL_MR_{variable_name}", f"LCL_MR_{variable_name}", f"MR_{variable_name}"],
            reference_column=f"MR_{variable_name}",
            custom_colors={f"m_UCL_MR_{variable_name}_AVG_1": "#d62728", f"m_LCL_MR_{variable_name}_AVG_1": "#d62728", f"m_MR_{variable_name}_AVG_1": "#ff7f0e"},
        )

    # === Dist plot ===
    custom_function = f"avg((1/(sqrt(2*PI())*${{std_{variable_name}}}))*" f"exp(-1*(pow({variable_name}-${{mean_{variable_name}}},2)/" f"(2*pow(${{std_{variable_name}}},2)))))"

    info = deepcopy(dist_plot_creation_info)
    chart_def = info["params"]["def"]

    # Insight Name
    chart_def["name"] = f"Dist Plot - {variable_name}"
    info["name"] = chart_def["name"]

    # X axis
    chart_def["genericDimension0"][0]["column"] = variable_name

    # Generate the custom function
    custom_function = f"avg((1/(sqrt(2*PI())*${{std_{variable_name}}}))*" f"exp(-1*(pow({variable_name}-${{mean_{variable_name}}},2)/" f"(2*pow(${{std_{variable_name}}},2)))))"

    # Generate generic measures dynamically from the template
    chart_def["genericMeasures"] = []
    for measure in deepcopy(dist_plot_generic_measures_template):
        updated = deepcopy(measure)
        if updated.get("column") == "__COLUMN__":
            updated["column"] = variable_name
        elif updated.get("column") == "__CUSTOM_COLUMN__":
            updated["column"] = f"norm_distribution_curve_{variable_name}"
        if "customFunction" in updated and updated["customFunction"] == "__CUSTOM_FUNCTION__":
            updated["customFunction"] = custom_function
        chart_def["genericMeasures"].append(updated)

    # Reference lines
    chart_def["referenceLines"][0]["datasetColumn"]["column"] = f"USL_{variable_name}"
    chart_def["referenceLines"][1]["datasetColumn"]["column"] = f"LSL_{variable_name}"

    # Target line

    # Remove existing Constant reference lines
    chart_def["referenceLines"] = [ref for ref in chart_def.get("referenceLines", []) if ref.get("sourceType") != "Constant"]

    # Check if target is defined in project variables
    target_key = f"TARGET_{variable_name}"
    has_target = target_key in variables["standard"]

    if has_target:
        target_value = variables["standard"][target_key]
        current_target_line = deepcopy(target_line)
        current_target_line["constantValue"] = target_value

        chart_def.setdefault("referenceLines", []).append(current_target_line)

    # Insight creation
    insights_ids[f"dist_{variable_name}_id"] = project.create_insight(info).insight_id

    if using_pa:
        # Make sure there are 3 variables in the chart, as PA rules always have min and max limits:
        while len(chart_def["genericMeasures"]) < 3:
            chart_def["genericMeasures"].append(deepcopy(generic_measure_template))
        # Chart name
        chart_def["name"] = f"{variable_name} time series"
        # Measures
        chart_def["genericMeasures"][0]["column"] = variable_name
        chart_def["genericMeasures"][1]["column"] = f"{variable_name}_selected_range_min"
        chart_def["genericMeasures"][2]["column"] = f"{variable_name}_selected_range_max"
        # Custom colors:
        chart_def["colorOptions"]["customColors"] = {
            f"m_{variable_name}_selected_range_min_AVG_1": "#9467bd",
            f"m_{variable_name}_AVG_1": "#030303",
            f"m_{variable_name}_selected_range_max_AVG_1": "#030303",
        }

    # === Metrics Insight ===

    # Get control limits based on chart type
    selected_df = df_by_subgroup
    if chart_type == "X_Bar":
        ucl_location_col = f"UCL_Xbar_{variable_name}"
        lcl_location_col = f"LCL_Xbar_{variable_name}"
        ucl_variability_col = f"UCL_R_{variable_name}"
        lcl_variability_col = f"LCL_R_{variable_name}"
    elif chart_type == "XS_Bar":
        ucl_location_col = f"UCL_Xbar_{variable_name}"
        lcl_location_col = f"LCL_Xbar_{variable_name}"
        ucl_variability_col = f"UCL_S_{variable_name}"
        lcl_variability_col = f"LCL_S_{variable_name}"
    elif chart_type == "I_MR":
        ucl_location_col = f"UCL_I_{variable_name}"
        lcl_location_col = f"LCL_I_{variable_name}"
        ucl_variability_col = f"UCL_MR_{variable_name}"
        lcl_variability_col = f"LCL_MR_{variable_name}"
        selected_df = df_raw

    if ucl_location_col in limits_row.columns and lcl_location_col in limits_row.columns and ucl_variability_col in limits_row.columns and lcl_variability_col in limits_row.columns:
        ucl_location_val = pd.to_numeric(limits_row.iloc[0][ucl_location_col], errors="coerce")
        lcl_location_val = pd.to_numeric(limits_row.iloc[0][lcl_location_col], errors="coerce")
        ucl_variability_val = pd.to_numeric(limits_row.iloc[0][ucl_variability_col], errors="coerce")
        lcl_variability_val = pd.to_numeric(limits_row.iloc[0][lcl_variability_col], errors="coerce")
    else:
        logger.warning(f"Missing UCL or LCL limits in selected_study_data_enhanced; skipping metrics for {variable_name}.")
        continue

    if pd.isna(ucl_location_val) or pd.isna(lcl_location_val) or pd.isna(ucl_variability_val) or pd.isna(lcl_variability_val):
        logger.warning(f"NaN UCL/LCL for {variable_name}; skipping metrics.")
        continue

    # Check control status
    under_control = is_under_control(
        selected_df,
        variable_name,
        ucl_location=float(ucl_location_val),
        lcl_location=float(lcl_location_val),
        ucl_variability=float(ucl_variability_val),
        lcl_variability=float(lcl_variability_val),
        chart_type=chart_type,
    )
    logger.info(f"For variable {variable_name}, under control status:{under_control}")
    if under_control[0]:
        current_metrics_creation_info = deepcopy(metrics_creation_info)
        current_metrics_creation_info["params"]["def"]["genericMeasures"] = []

        # Add Cp and Cpk if available
        if f"{variable_name}_cp" in available_capability_metrics and f"{variable_name}_cpk" in available_capability_metrics:
            current_metrics_creation_info["params"]["def"]["genericMeasures"].append(make_metric_block(f"{variable_name}_cp", "cp"))
            current_metrics_creation_info["params"]["def"]["genericMeasures"].append(make_metric_block(f"{variable_name}_cpk", "cpk"))

        # Add Cpm and Cpkm if target is defined and metrics available
        if has_target and f"{variable_name}_cpm" in available_capability_metrics and f"{variable_name}_cpkm" in available_capability_metrics:
            current_metrics_creation_info["params"]["def"]["genericMeasures"].append(make_metric_block(f"{variable_name}_cpm", "cpm"))
            current_metrics_creation_info["params"]["def"]["genericMeasures"].append(make_metric_block(f"{variable_name}_cpkm", "cpkm"))

        # Only create the insight if any metric was added
        if current_metrics_creation_info["params"]["def"]["genericMeasures"]:
            current_metrics_creation_info["params"]["def"]["name"] = f"{variable_name}_metrics"
            current_metrics_creation_info["name"] = f"Metrics - {variable_name}"
            insights_ids[f"metrics_{variable_name}_id"] = project.create_insight(current_metrics_creation_info).insight_id


# Creating cpk_table
cpk_table_creation_info = getattr(insights_static, "cpk_table_creation_info")
insights_ids["cpk_table_id"] = project.create_insight(cpk_table_creation_info).insight_id

# Save insights_ids in variables:
variables["standard"]["insights_ids"] = insights_ids
project.set_variables(variables)

# Build dashboard

client = dataiku.api_client()
project, variables = get_current_project_and_variables()

selected_variables = variables["standard"]["selected_parameterz"]
insights_ids = variables["standard"]["insights_ids"]

dashboard = project.get_dashboard("1BAutEI")
dashboard_settings = dashboard.get_settings()
dashboard_raw_settings = dashboard_settings.get_raw()


# === Configuring Control charts (1st tab) ===
second_page_tiles = dashboard_raw_settings["pages"][0]["grid"]["tiles"]

# Keep the initial control charts intro tile
intro_tile = None
for tile in second_page_tiles:
    if tile.get("tileType") == "TEXT" and "Control charts are statistical tools" in tile.get("tileParams", {}).get("text", ""):
        intro_tile = tile
        break

# Build new layout starting with the intro
new_tiles = [intro_tile] if intro_tile else []
offset = 3  # Start below the intro

for var in selected_variables:
    # Determine chart titles based on chart_type
    if chart_type == "X_Bar":
        location_chart_title = f"X-bar - {var}"
        variability_chart_title = f"R - {var}"
        text = f"## {var}\nSubgroup size: {subgroup_size}"
    elif chart_type == "XS_Bar":
        location_chart_title = f"X-bar - {var}"
        variability_chart_title = f"S - {var}"
        text = f"## {var}\nSubgroup size: {subgroup_size}"
    elif chart_type == "I_MR":
        location_chart_title = f"I (Individual) - {var}"
        variability_chart_title = f"MR (Moving Range) - {var}"
        text = f"## {var}"

    # Section header
    new_tiles.append(
        {
            "tileType": "TEXT",
            "box": {"top": offset, "left": 0, "width": 36, "height": 2},
            "tileParams": {"text": text, "textAlign": "LEFT", "verticalAlign": "TOP"},
            "backgroundColor": "#ffffff",
            "backgroundOpacity": 1.0,
            "borderOptions": {"color": "#D9D9D9", "radius": 0, "size": 1},
            "titleOptions": {"showTitle": "NO", "fontColor": "#000", "fontSize": 13},
            "autoLoad": True,
            "locked": False,
            "displayMode": "INSIGHT",
            "useDashboardSpacing": True,
            "tileSpacing": 8,
            "resizeImageMode": "FIT_SIZE",
            "clickAction": "DO_NOTHING",
        }
    )

    # Location chart (X-bar or I)
    new_tiles.append(
        {
            "tileType": "INSIGHT",
            "box": {"top": offset + 2, "left": 0, "width": 18, "height": 9},
            "insightId": insights_ids[f"x_bar_{var}_id"],
            "insightType": "chart",
            "tileParams": {
                "showXAxis": True,
                "showYAxis": True,
                "showLegend": True,
                "showBrush": False,
                "inheritLegendPlacement": True,
                "legendPlacement": "OUTER_RIGHT",
                "showTooltips": True,
                "autoPlayAnimation": True,
                "useInsightTheme": False,
            },
            "backgroundColor": "#ffffff",
            "backgroundOpacity": 1.0,
            "borderOptions": {"color": "#D9D9D9", "radius": 0, "size": 1},
            "titleOptions": {"showTitle": "YES", "title": location_chart_title, "fontColor": "#000", "fontSize": 13},
            "autoLoad": True,
            "locked": False,
            "displayMode": "INSIGHT",
            "useDashboardSpacing": True,
            "tileSpacing": 8,
            "resizeImageMode": "FIT_SIZE",
            "clickAction": "DO_NOTHING",
        }
    )

    # Variability chart (R, S, or MR)
    new_tiles.append(
        {
            "tileType": "INSIGHT",
            "box": {"top": offset + 2, "left": 18, "width": 18, "height": 9},
            "insightId": insights_ids[f"variability_{var}_id"],
            "insightType": "chart",
            "tileParams": {
                "showXAxis": True,
                "showYAxis": True,
                "showLegend": True,
                "showBrush": False,
                "inheritLegendPlacement": True,
                "legendPlacement": "OUTER_RIGHT",
                "showTooltips": True,
                "autoPlayAnimation": True,
                "useInsightTheme": False,
            },
            "backgroundColor": "#ffffff",
            "backgroundOpacity": 1.0,
            "borderOptions": {"color": "#D9D9D9", "radius": 0, "size": 1},
            "titleOptions": {"showTitle": "YES", "title": variability_chart_title, "fontColor": "#000", "fontSize": 13},
            "autoLoad": True,
            "locked": False,
            "displayMode": "INSIGHT",
            "useDashboardSpacing": True,
            "tileSpacing": 8,
            "resizeImageMode": "FIT_SIZE",
            "clickAction": "DO_NOTHING",
        }
    )

    offset += 11

dashboard_raw_settings["pages"][0]["grid"]["tiles"] = new_tiles
dashboard_settings.save()

# === Configuring Capability Study (2nd tab) ===
third_page_tiles = dashboard_raw_settings["pages"][1]["grid"]["tiles"]

# --- Keep the initial capability study intro tile ---
capability_intro_tile = None
for tile in third_page_tiles:
    if tile.get("tileType") == "TEXT" and "This page contains:" in tile.get("tileParams", {}).get("text", ""):
        capability_intro_tile = tile
        break

# --- Start new layout ---
new_capability_tiles = [capability_intro_tile] if capability_intro_tile else []

offset = 5  # Start after the intro tile

for i, var in enumerate(selected_variables):
    var_offset = offset + i * 10

    # Section title
    new_capability_tiles.append(
        {
            "tileType": "TEXT",
            "box": {"top": var_offset - 2, "left": 0, "width": 36, "height": 2},
            "tileParams": {"text": f"## {var}", "textAlign": "LEFT", "verticalAlign": "TOP"},
            "backgroundColor": "#ffffff",
            "backgroundOpacity": 1.0,
            "borderOptions": {"color": "#D9D9D9", "radius": 0, "size": 1},
            "titleOptions": {"showTitle": "NO", "fontColor": "#000", "fontSize": 13},
            "autoLoad": True,
            "locked": False,
            "displayMode": "INSIGHT",
            "useDashboardSpacing": True,
            "tileSpacing": 8,
            "resizeImageMode": "FIT_SIZE",
            "clickAction": "DO_NOTHING",
        }
    )

    # Distribution plot
    new_capability_tiles.append(
        {
            "tileType": "INSIGHT",
            "box": {"top": var_offset, "left": 3, "width": 21, "height": 8},
            "insightId": insights_ids[f"dist_{var}_id"],
            "insightType": "chart",
            "tileParams": {
                "showXAxis": True,
                "showYAxis": True,
                "showLegend": True,
                "showBrush": False,
                "inheritLegendPlacement": True,
                "legendPlacement": "OUTER_RIGHT",
                "showTooltips": True,
                "autoPlayAnimation": True,
                "useInsightTheme": False,
            },
            "backgroundColor": "#ffffff",
            "backgroundOpacity": 1.0,
            "borderOptions": {"color": "#D9D9D9", "radius": 0, "size": 1},
            "titleOptions": {"showTitle": "NO", "title": f"Specification limits - {var}", "fontColor": "#000", "fontSize": 13},
            "autoLoad": True,
            "locked": False,
            "displayMode": "INSIGHT",
            "useDashboardSpacing": True,
            "tileSpacing": 8,
            "resizeImageMode": "FIT_SIZE",
            "clickAction": "DO_NOTHING",
        }
    )

    # Dataset view (CPK table)
    new_capability_tiles.append(
        {
            "tileType": "INSIGHT",
            "box": {"top": var_offset, "left": 24, "width": 12, "height": 8},
            "insightId": insights_ids["cpk_table_id"],
            "insightType": "dataset_table",
            "tileParams": {"viewKind": "EXPLORE", "showName": True, "showDescription": True, "showCustomFields": True, "showStorageType": False, "showMeaning": False, "showProgressBar": False},
            "backgroundColor": "#ffffff",
            "backgroundOpacity": 1.0,
            "borderOptions": {"color": "#CCCCCC", "radius": 0, "size": 1},
            "titleOptions": {"showTitle": "YES", "title": "CPK Meanings", "fontColor": "#000000", "fontSize": 13},
            "autoLoad": True,
            "locked": False,
            "displayMode": "INSIGHT",
            "useDashboardSpacing": True,
            "tileSpacing": 8,
            "resizeImageMode": "FIT_SIZE",
            "clickAction": "DO_NOTHING",
        }
    )

    # Add metrics insight if available, otherwise fallback to explanatory text
    metrics_id = insights_ids.get(f"metrics_{var}_id")
    if metrics_id:
        new_capability_tiles.append(
            {
                "tileType": "INSIGHT",
                "box": {"top": var_offset, "left": 0, "width": 3, "height": 8},
                "insightId": metrics_id,
                "insightType": "chart",
                "tileParams": {
                    "showXAxis": True,
                    "showYAxis": True,
                    "showLegend": False,
                    "showBrush": False,
                    "inheritLegendPlacement": True,
                    "legendPlacement": "OUTER_RIGHT",
                    "showTooltips": True,
                    "autoPlayAnimation": True,
                    "useInsightTheme": False,
                },
                "backgroundColor": "#ffffff",
                "backgroundOpacity": 1.0,
                "borderOptions": {"color": "#D9D9D9", "radius": 0, "size": 1},
                "titleOptions": {"showTitle": "NO", "fontColor": "#000", "fontSize": 13},
                "autoLoad": True,
                "locked": False,
                "displayMode": "INSIGHT",
                "useDashboardSpacing": True,
                "tileSpacing": 8,
                "resizeImageMode": "FIT_SIZE",
                "clickAction": "DO_NOTHING",
            }
        )
    else:
        # Check if specification limits exist in project variables
        USL = variables["standard"].get(f"USL_{var}", None)
        LSL = variables["standard"].get(f"LSL_{var}", None)

        if USL is None or LSL is None:
            message_text = f"_{var}_   **has no specification limits (USL/LSL) defined, so capability metrics are not calculated**"
        else:
            message_text = f"_{var}_   **is not under statistical control, so capability metrics are not calculated**"

        new_capability_tiles.append(
            {
                "tileType": "TEXT",
                "box": {"top": var_offset, "left": 0, "width": 3, "height": 8},
                "clickAction": "DO_NOTHING",
                "tileParams": {"text": message_text, "textAlign": "LEFT", "verticalAlign": "TOP"},
                "backgroundOpacity": 1.0,
                "backgroundColor": "#ffffff",
                "autoLoad": True,
                "locked": False,
                "borderOptions": {"color": "#D9D9D9", "radius": 0, "size": 1},
                "titleOptions": {"showTitle": "NO", "fontColor": "#000", "fontSize": 13},
                "displayMode": "INSIGHT",
                "useDashboardSpacing": True,
                "tileSpacing": 8,
                "resizeImageMode": "FIT_SIZE",
            }
        )

dashboard_raw_settings["pages"][1]["grid"]["tiles"] = new_capability_tiles
dashboard_settings.save()
