import dataiku
from parameters_control_utils import generate_process_capability_visualization_data, generate_selected_study_data_enhanced

# Load input dataset
input_dku_dataset = dataiku.Dataset("experiment_dataset_sorted")
df_raw = input_dku_dataset.get_dataframe()

# Get measurement columns and chart_type from project variables
client = dataiku.api_client()
project = client.get_default_project()
project_vars = project.get_variables()["standard"]
measurement_columns = project_vars.get("selected_parameterz", [])
chart_type = project_vars.get("chart_type", "X_Bar")

if not measurement_columns:
    raise ValueError("No 'selected_parameters' found in project variables. Please define them to proceed with analysis.")

# 1. Generate data for Process Capability Visualization
process_capability_viz_df = generate_process_capability_visualization_data(df_raw, measurement_columns, chart_type)
process_capability_visualization_output = dataiku.Dataset("process_capability_visualization")
process_capability_visualization_output.write_with_schema(process_capability_viz_df)

# 2. Generate Enhanced Selected Study Data
selected_study_data_enhanced_df = generate_selected_study_data_enhanced(df_raw, measurement_columns, chart_type)
selected_study_data_enhanced_output = dataiku.Dataset("selected_study_data_enhanced")
selected_study_data_enhanced_output.write_with_schema(selected_study_data_enhanced_df)
