# -*- coding: utf-8 -*-
import dataiku
import pandas as pd, numpy as np
from dataiku import pandasutils as pdu
import ast

# Get variables:
project = dataiku.api_client().get_default_project()
variables = project.get_variables()

# Read recipe inputs
input_pa_project_key = variables["standard"]["pa_instance"]
data_source_mode = variables["standard"]["data_source_mode"]

if data_source_mode == "file_upload":
    # Dataset saved_studies renamed to saved_studies-1 by romain.menini@dataiku.com on 2025-09-23 08:33:12
    input_dataset_saved_studies = dataiku.Dataset("saved_studies")
else:
    input_dataset_saved_studies = dataiku.Dataset(f"{input_pa_project_key}.saved_studies")

df = input_dataset_saved_studies.get_dataframe(infer_with_pandas=False)


# To be updated by user's selection
study_id = variables["standard"]["saved_study"]
row = df.loc[df['study_id'] == study_id].iloc[0]

# Extract variables to save as project variables
source_project_key = row['source_project_key']
source_dataset= row['source_datasets']
source_dataset = source_dataset.strip("[]'\"")
date_column = row['date_column']
date_range = row['date_range']

# Set project variables

variables["standard"]["source_project_key"] = source_project_key
variables["standard"]["source_dataset"] = source_dataset
variables["local"]["date_column"] = date_column
variables["standard"]["date_range"] = date_range
project.set_variables(variables)

# Parse the string representations into Python objects
target_info = ast.literal_eval(row['target_info'])
study_variables = ast.literal_eval(row['variables'])

# Prepare the target_info and analysis_variables for output dataset
output_data = {
    'target': target_info['target'],
    'target_date_from': target_info['target_date']['from'],
    'target_date_to': target_info['target_date']['to'],
    'target_selected_range_min': target_info['target_selected_range'][0],
    'target_selected_range_max': target_info['target_selected_range'][1],
    'variables': study_variables,
    'selected_variables': [item['variable'] for item in study_variables]

}

# Convert output_data to DataFrame
output_df = pd.DataFrame([output_data])

# Write recipe outputs
saved_study_extracted = dataiku.Dataset("saved_study_extracted")
saved_study_extracted.write_with_schema(output_df)