import json
import sys
import os
from io import StringIO
import traceback

cwd = os.getcwd()
files = sorted(os.listdir('.'))
print(f'Working directory: {cwd}')
print(f'Files: {files}')

notebooks = [
    'research_asset_class_momentum.ipynb',
    'research_volatility_regime_ml.ipynb',
    'research_defensive_etf_rotation.ipynb',
    'research_commodity_term_structure.ipynb',
    'research_piotroski_fscore.ipynb',
    'research_long_short_harvest.ipynb',
    'research_macro_factor_rotation.ipynb',
    'research_puppies_of_dow.ipynb',
]

for nb_name in notebooks:
    print(f'\n{"=" * 60}')
    print(f'Processing: {nb_name}')
    try:
        with open(nb_name, 'r') as f:
            nb = json.load(f)
    except FileNotFoundError:
        print(f'SKIP: {nb_name} not found')
        continue
    exec_count = 0
    success = 0
    errors = 0
    for cell in nb['cells']:
        if cell['cell_type'] != 'code':
            continue
        exec_count += 1
        code = ''.join(cell['source'])
        old_stdout, old_stderr = sys.stdout, sys.stderr
        sys.stdout, sys.stderr = StringIO(), StringIO()
        try:
            exec(code, globals())
            cell['execution_count'] = exec_count
            out = sys.stdout.getvalue()
            err = sys.stderr.getvalue()
            outputs = []
            if out:
                lines = out.splitlines(True)
                if lines and not lines[-1].endswith('\n'):
                    lines[-1] += '\n'
                outputs.append({'output_type': 'stream', 'name': 'stdout', 'text': lines})
            if err:
                outputs.append({'output_type': 'stream', 'name': 'stderr', 'text': err.splitlines(True)})
            cell['outputs'] = outputs
            success += 1
        except Exception as e:
            cell['execution_count'] = exec_count
            cell['outputs'] = [{
                'output_type': 'error',
                'ename': type(e).__name__,
                'evalue': str(e),
                'traceback': [traceback.format_exc()]
            }]
            errors += 1
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr
    out_name = nb_name.replace('.ipynb', '_executed.ipynb')
    with open(out_name, 'w') as f:
        json.dump(nb, f, indent=1)
    print(f'  {success} OK, {errors} ERR -> {out_name}')
print('\n\nALL DONE')