src/data_collection/collect_entsoe.py.backup

"""ENTSO-E Transparency Platform Data Collection with Rate Limiting

Collects generation, load, and cross-border flow data from ENTSO-E API.
Implements proper rate limiting to avoid temporary bans.

ENTSO-E Rate Limits (OFFICIAL):
- 60 requests per 60 seconds (hard limit - exceeding triggers 10-min ban)
- Screen scraping >60 requests/min leads to temporary IP ban

Strategy:
- 27 requests/minute (45% of 60 limit - safe)
- 1 request every ~2.2 seconds
- Request data in monthly chunks to minimize API calls
"""

import polars as pl
from pathlib import Path
from datetime import datetime, timedelta
from dotenv import load_dotenv
import os
import time
from typing import List, Tuple
from tqdm import tqdm
from entsoe import EntsoePandasClient
import pandas as pd
import zipfile
from io import BytesIO
import xml.etree.ElementTree as ET


# Load environment variables
load_dotenv()


# FBMC Bidding Zones (12 zones from project plan)
BIDDING_ZONES = {
    'AT': 'Austria',
    'BE': 'Belgium',
    'HR': 'Croatia',
    'CZ': 'Czech Republic',
    'FR': 'France',
    'DE_LU': 'Germany-Luxembourg',
    'HU': 'Hungary',
    'NL': 'Netherlands',
    'PL': 'Poland',
    'RO': 'Romania',
    'SK': 'Slovakia',
    'SI': 'Slovenia',
}


# FBMC Cross-Border Flows (~20 major borders)
BORDERS = [
    ('DE_LU', 'NL'),
    ('DE_LU', 'FR'),
    ('DE_LU', 'BE'),
    ('DE_LU', 'AT'),
    ('DE_LU', 'CZ'),
    ('DE_LU', 'PL'),
    ('FR', 'BE'),
    ('FR', 'ES'),  # External but affects FBMC
    ('FR', 'CH'),  # External but affects FBMC
    ('AT', 'CZ'),
    ('AT', 'HU'),
    ('AT', 'SI'),
    ('AT', 'CH'),  # External but affects FBMC
    ('CZ', 'SK'),
    ('CZ', 'PL'),
    ('HU', 'SK'),
    ('HU', 'RO'),
    ('HU', 'HR'),
    ('SI', 'HR'),
    ('PL', 'SK'),
    ('PL', 'CZ'),
]


# FBMC Bidding Zone EIC Codes (for asset-specific outages)
BIDDING_ZONE_EICS = {
    'AT': '10YAT-APG------L',
    'BE': '10YBE----------2',
    'HR': '10YHR-HEP------M',
    'CZ': '10YCZ-CEPS-----N',
    'FR': '10YFR-RTE------C',
    'DE_LU': '10Y1001A1001A82H',
    'HU': '10YHU-MAVIR----U',
    'NL': '10YNL----------L',
    'PL': '10YPL-AREA-----S',
    'RO': '10YRO-TEL------P',
    'SK': '10YSK-SEPS-----K',
    'SI': '10YSI-ELES-----O',
    'CH': '10YCH-SWISSGRIDZ',
}


# PSR Types for generation data collection
PSR_TYPES = {
    'B01': 'Biomass',
    'B02': 'Fossil Brown coal/Lignite',
    'B03': 'Fossil Coal-derived gas',
    'B04': 'Fossil Gas',
    'B05': 'Fossil Hard coal',
    'B06': 'Fossil Oil',
    'B07': 'Fossil Oil shale',
    'B08': 'Fossil Peat',
    'B09': 'Geothermal',
    'B10': 'Hydro Pumped Storage',
    'B11': 'Hydro Run-of-river and poundage',
    'B12': 'Hydro Water Reservoir',
    'B13': 'Marine',
    'B14': 'Nuclear',
    'B15': 'Other renewable',
    'B16': 'Solar',
    'B17': 'Waste',
    'B18': 'Wind Offshore',
    'B19': 'Wind Onshore',
    'B20': 'Other',
}


# Zones with significant pumped storage capacity
PUMPED_STORAGE_ZONES = ['CH', 'AT', 'DE_LU', 'FR', 'HU', 'PL', 'RO']


# Zones with significant hydro reservoir capacity
HYDRO_RESERVOIR_ZONES = ['CH', 'AT', 'FR', 'RO', 'SI', 'HR', 'SK']


# Zones with nuclear generation
NUCLEAR_ZONES = ['FR', 'BE', 'CZ', 'HU', 'RO', 'SI', 'SK']


class EntsoECollector:
    """Collect ENTSO-E data with proper rate limiting."""

    def __init__(self, requests_per_minute: int = 27):
        """Initialize collector with rate limiting.

        Args:
            requests_per_minute: Max requests per minute (default: 27 = 45% of 60 limit)
        """
        api_key = os.getenv('ENTSOE_API_KEY')
        if not api_key or 'your_entsoe' in api_key.lower():
            raise ValueError("ENTSO-E API key not configured in .env file")

        self.client = EntsoePandasClient(api_key=api_key)
        self.requests_per_minute = requests_per_minute
        self.delay_seconds = 60.0 / requests_per_minute
        self.request_count = 0

        print(f"ENTSO-E Collector initialized")
        print(f"Rate limit: {self.requests_per_minute} requests/minute")
        print(f"Delay between requests: {self.delay_seconds:.2f}s")

    def _rate_limit(self):
        """Apply rate limiting delay."""
        time.sleep(self.delay_seconds)
        self.request_count += 1

    def _generate_monthly_chunks(
        self,
        start_date: str,
        end_date: str
    ) -> List[Tuple[pd.Timestamp, pd.Timestamp]]:
        """Generate yearly date chunks for API requests (OPTIMIZED).

        ENTSO-E API supports up to 1 year per request, so we use yearly chunks
        instead of monthly to reduce API calls by 12x.

        Args:
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            List of (start, end) timestamp tuples
        """
        start_dt = pd.Timestamp(start_date, tz='UTC')
        end_dt = pd.Timestamp(end_date, tz='UTC')

        chunks = []
        current = start_dt

        while current < end_dt:
            # Get end of year or end_date, whichever is earlier
            year_end = pd.Timestamp(f"{current.year}-12-31 23:59:59", tz='UTC')
            chunk_end = min(year_end, end_dt)

            chunks.append((current, chunk_end))
            current = chunk_end + pd.Timedelta(hours=1)

        return chunks

    def collect_generation_per_type(
        self,
        zone: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect generation by production type for a bidding zone.

        Args:
            zone: Bidding zone code (e.g., 'DE_LU', 'FR')
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with generation data
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} generation", leave=False):
            try:
                # Fetch generation data
                df = self.client.query_generation(
                    zone,
                    start=start_chunk,
                    end=end_chunk,
                    psr_type=None  # Get all production types
                )

                if df is not None and not df.empty:
                    # Convert to long format
                    df_reset = df.reset_index()
                    df_melted = df_reset.melt(
                        id_vars=['index'],
                        var_name='production_type',
                        value_name='generation_mw'
                    )
                    df_melted = df_melted.rename(columns={'index': 'timestamp'})
                    df_melted['zone'] = zone

                    # Convert to Polars
                    pl_df = pl.from_pandas(df_melted)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    ❌ Failed {zone} {start_chunk.date()} to {end_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_load(
        self,
        zone: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect load (demand) data for a bidding zone.

        Args:
            zone: Bidding zone code
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with load data
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} load", leave=False):
            try:
                # Fetch load data
                series = self.client.query_load(
                    zone,
                    start=start_chunk,
                    end=end_chunk
                )

                if series is not None and not series.empty:
                    df = pd.DataFrame({
                        'timestamp': series.index,
                        'load_mw': series.values,
                        'zone': zone
                    })

                    pl_df = pl.from_pandas(df)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    ❌ Failed {zone} {start_chunk.date()} to {end_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_cross_border_flows(
        self,
        from_zone: str,
        to_zone: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect cross-border flow data between two zones.

        Args:
            from_zone: From bidding zone
            to_zone: To bidding zone
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with flow data
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        border_id = f"{from_zone}_{to_zone}"

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {border_id}", leave=False):
            try:
                # Fetch cross-border flow
                series = self.client.query_crossborder_flows(
                    from_zone,
                    to_zone,
                    start=start_chunk,
                    end=end_chunk
                )

                if series is not None and not series.empty:
                    df = pd.DataFrame({
                        'timestamp': series.index,
                        'flow_mw': series.values,
                        'from_zone': from_zone,
                        'to_zone': to_zone,
                        'border': border_id
                    })

                    pl_df = pl.from_pandas(df)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    ❌ Failed {border_id} {start_chunk.date()} to {end_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_transmission_outages_asset_specific(
        self,
        cnec_eics: List[str],
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect asset-specific transmission outages using XML parsing.

        Uses validated Phase 1C/1D methodology: Query border-level outages,
        parse ZIP/XML to extract Asset_RegisteredResource.mRID elements,
        filter to CNEC EIC codes.

        Args:
            cnec_eics: List of CNEC EIC codes to filter (e.g., 200 critical CNECs)
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with outage events
            Columns: asset_eic, asset_name, start_time, end_time,
                     businesstype, from_zone, to_zone, border
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_outages = []

        # Query all FBMC borders for transmission outages
        for zone1, zone2 in tqdm(BORDERS, desc="Transmission outages (borders)"):
            zone1_eic = BIDDING_ZONE_EICS.get(zone1)
            zone2_eic = BIDDING_ZONE_EICS.get(zone2)

            if not zone1_eic or not zone2_eic:
                continue

            for start_chunk, end_chunk in chunks:
                try:
                    # Query border-level outages (raw bytes)
                    response = self.client._base_request(
                        params={
                            'documentType': 'A78',  # Transmission unavailability
                            'in_Domain': zone2_eic,
                            'out_Domain': zone1_eic
                        },
                        start=start_chunk,
                        end=end_chunk
                    )

                    outages_zip = response.content

                    # Parse ZIP and extract Asset_RegisteredResource.mRID
                    with zipfile.ZipFile(BytesIO(outages_zip), 'r') as zf:
                        xml_files = [f for f in zf.namelist() if f.endswith('.xml')]

                        for xml_file in xml_files:
                            with zf.open(xml_file) as xf:
                                xml_content = xf.read()
                                root = ET.fromstring(xml_content)

                                # Get namespace
                                nsmap = dict([node for _, node in ET.iterparse(
                                    BytesIO(xml_content), events=['start-ns']
                                )])
                                ns_uri = nsmap.get('', None)

                                # Find TimeSeries elements
                                if ns_uri:
                                    timeseries_found = root.findall('.//{' + ns_uri + '}TimeSeries')
                                else:
                                    timeseries_found = root.findall('.//TimeSeries')

                                for ts in timeseries_found:
                                    # Extract Asset_RegisteredResource.mRID
                                    if ns_uri:
                                        reg_resource = ts.find('.//{' + ns_uri + '}Asset_RegisteredResource')
                                    else:
                                        reg_resource = ts.find('.//Asset_RegisteredResource')

                                    if reg_resource is not None:
                                        # Get asset EIC
                                        if ns_uri:
                                            mrid_elem = reg_resource.find('.//{' + ns_uri + '}mRID')
                                            name_elem = reg_resource.find('.//{' + ns_uri + '}name')
                                        else:
                                            mrid_elem = reg_resource.find('.//mRID')
                                            name_elem = reg_resource.find('.//name')

                                        if mrid_elem is not None:
                                            asset_eic = mrid_elem.text

                                            # Filter to CNEC list
                                            if asset_eic in cnec_eics:
                                                asset_name = name_elem.text if name_elem is not None else ''

                                                # Extract outage periods
                                                if ns_uri:
                                                    periods = ts.findall('.//{' + ns_uri + '}Available_Period')
                                                else:
                                                    periods = ts.findall('.//Available_Period')

                                                for period in periods:
                                                    if ns_uri:
                                                        time_interval = period.find('.//{' + ns_uri + '}timeInterval')
                                                    else:
                                                        time_interval = period.find('.//timeInterval')

                                                    if time_interval is not None:
                                                        if ns_uri:
                                                            start_elem = time_interval.find('.//{' + ns_uri + '}start')
                                                            end_elem = time_interval.find('.//{' + ns_uri + '}end')
                                                        else:
                                                            start_elem = time_interval.find('.//start')
                                                            end_elem = time_interval.find('.//end')

                                                        if start_elem is not None and end_elem is not None:
                                                            # Extract business type from root
                                                            if ns_uri:
                                                                business_type_elem = root.find('.//{' + ns_uri + '}businessType')
                                                            else:
                                                                business_type_elem = root.find('.//businessType')

                                                            business_type = business_type_elem.text if business_type_elem is not None else 'Unknown'

                                                            all_outages.append({
                                                                'asset_eic': asset_eic,
                                                                'asset_name': asset_name,
                                                                'start_time': pd.Timestamp(start_elem.text),
                                                                'end_time': pd.Timestamp(end_elem.text),
                                                                'businesstype': business_type,
                                                                'from_zone': zone1,
                                                                'to_zone': zone2,
                                                                'border': f"{zone1}_{zone2}"
                                                            })

                    self._rate_limit()

                except Exception as e:
                    # Empty response or no outages is OK
                    if "empty" not in str(e).lower():
                        print(f"    Warning: {zone1}->{zone2} {start_chunk.date()}: {e}")
                    self._rate_limit()
                    continue

        if all_outages:
            return pl.DataFrame(all_outages)
        else:
            return pl.DataFrame()

    def collect_day_ahead_prices(
        self,
        zone: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect day-ahead electricity prices.

        Args:
            zone: Bidding zone code
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with price data
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} prices", leave=False):
            try:
                series = self.client.query_day_ahead_prices(
                    zone,
                    start=start_chunk,
                    end=end_chunk
                )

                if series is not None and not series.empty:
                    df = pd.DataFrame({
                        'timestamp': series.index,
                        'price_eur_mwh': series.values,
                        'zone': zone
                    })

                    pl_df = pl.from_pandas(df)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    Warning: {zone} {start_chunk.date()} to {end_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_hydro_reservoir_storage(
        self,
        zone: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect hydro reservoir storage levels (weekly data).

        Args:
            zone: Bidding zone code
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with reservoir storage data (weekly)
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} hydro storage", leave=False):
            try:
                series = self.client.query_aggregate_water_reservoirs_and_hydro_storage(
                    zone,
                    start=start_chunk,
                    end=end_chunk
                )

                if series is not None and not series.empty:
                    df = pd.DataFrame({
                        'timestamp': series.index,
                        'storage_mwh': series.values,
                        'zone': zone
                    })

                    pl_df = pl.from_pandas(df)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    Warning: {zone} {start_chunk.date()} to {end_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_pumped_storage_generation(
        self,
        zone: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect pumped storage generation (B10 PSR type).

        Note: Consumption data not separately available from ENTSO-E API.
        Returns generation-only data.

        Args:
            zone: Bidding zone code
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with pumped storage generation
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} pumped storage", leave=False):
            try:
                series = self.client.query_generation(
                    zone,
                    start=start_chunk,
                    end=end_chunk,
                    psr_type='B10'  # Hydro Pumped Storage
                )

                if series is not None and not series.empty:
                    # Handle both Series and DataFrame returns
                    if isinstance(series, pd.DataFrame):
                        # If multiple columns, take first
                        series = series.iloc[:, 0]

                    df = pd.DataFrame({
                        'timestamp': series.index,
                        'generation_mw': series.values,
                        'zone': zone
                    })

                    pl_df = pl.from_pandas(df)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    Warning: {zone} {start_chunk.date()} to {end_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_load_forecast(
        self,
        zone: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect load forecast data.

        Args:
            zone: Bidding zone code
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with load forecast
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} load forecast", leave=False):
            try:
                series = self.client.query_load_forecast(
                    zone,
                    start=start_chunk,
                    end=end_chunk
                )

                if series is not None and not series.empty:
                    df = pd.DataFrame({
                        'timestamp': series.index,
                        'forecast_mw': series.values,
                        'zone': zone
                    })

                    pl_df = pl.from_pandas(df)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    Warning: {zone} {start_chunk.date()} to {end_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_generation_outages(
        self,
        zone: str,
        start_date: str,
        end_date: str,
        psr_type: str = None
    ) -> pl.DataFrame:
        """Collect generation/production unit outages.

        Uses document type A77 (unavailability of generation units).
        Particularly important for nuclear planned outages which are known
        months in advance and significantly impact cross-border flows.

        Args:
            zone: Bidding zone code
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)
            psr_type: Optional PSR type filter (B14=Nuclear, B04=Gas, B05=Coal, etc.)

        Returns:
            Polars DataFrame with generation unit outages
            Columns: unit_name, psr_type, psr_name, capacity_mw,
                     start_time, end_time, businesstype, zone
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_outages = []

        zone_eic = BIDDING_ZONE_EICS.get(zone)
        if not zone_eic:
            return pl.DataFrame()

        psr_name = PSR_TYPES.get(psr_type, psr_type) if psr_type else 'All'

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} {psr_name} outages", leave=False):
            try:
                # Build query parameters
                params = {
                    'documentType': 'A77',  # Generation unavailability
                    'biddingZone_Domain': zone_eic
                }

                # Add PSR type filter if specified
                if psr_type:
                    params['psrType'] = psr_type

                # Query generation unavailability
                response = self.client._base_request(
                    params=params,
                    start=start_chunk,
                    end=end_chunk
                )

                outages_zip = response.content

                # Parse ZIP and extract outage information
                with zipfile.ZipFile(BytesIO(outages_zip), 'r') as zf:
                    xml_files = [f for f in zf.namelist() if f.endswith('.xml')]

                    for xml_file in xml_files:
                        with zf.open(xml_file) as xf:
                            xml_content = xf.read()
                            root = ET.fromstring(xml_content)

                            # Get namespace
                            nsmap = dict([node for _, node in ET.iterparse(
                                BytesIO(xml_content), events=['start-ns']
                            )])
                            ns_uri = nsmap.get('', None)

                            # Find TimeSeries elements
                            if ns_uri:
                                timeseries_found = root.findall('.//{' + ns_uri + '}TimeSeries')
                            else:
                                timeseries_found = root.findall('.//TimeSeries')

                            for ts in timeseries_found:
                                # Extract production unit information
                                if ns_uri:
                                    prod_unit = ts.find('.//{' + ns_uri + '}Production_RegisteredResource')
                                else:
                                    prod_unit = ts.find('.//Production_RegisteredResource')

                                if prod_unit is not None:
                                    # Get unit details
                                    if ns_uri:
                                        name_elem = prod_unit.find('.//{' + ns_uri + '}name')
                                        psr_elem = prod_unit.find('.//{' + ns_uri + '}psrType')
                                    else:
                                        name_elem = prod_unit.find('.//name')
                                        psr_elem = prod_unit.find('.//psrType')

                                    unit_name = name_elem.text if name_elem is not None else 'Unknown'
                                    unit_psr = psr_elem.text if psr_elem is not None else psr_type

                                    # Extract outage periods and capacity
                                    if ns_uri:
                                        periods = ts.findall('.//{' + ns_uri + '}Unavailable_Period')
                                    else:
                                        periods = ts.findall('.//Unavailable_Period')

                                    for period in periods:
                                        if ns_uri:
                                            time_interval = period.find('.//{' + ns_uri + '}timeInterval')
                                            quantity_elem = period.find('.//{' + ns_uri + '}quantity')
                                        else:
                                            time_interval = period.find('.//timeInterval')
                                            quantity_elem = period.find('.//quantity')

                                        if time_interval is not None:
                                            if ns_uri:
                                                start_elem = time_interval.find('.//{' + ns_uri + '}start')
                                                end_elem = time_interval.find('.//{' + ns_uri + '}end')
                                            else:
                                                start_elem = time_interval.find('.//start')
                                                end_elem = time_interval.find('.//end')

                                            if start_elem is not None and end_elem is not None:
                                                # Get business type
                                                if ns_uri:
                                                    business_type_elem = root.find('.//{' + ns_uri + '}businessType')
                                                else:
                                                    business_type_elem = root.find('.//businessType')

                                                business_type = business_type_elem.text if business_type_elem is not None else 'Unknown'

                                                # Get capacity
                                                capacity_mw = float(quantity_elem.text) if quantity_elem is not None else 0.0

                                                all_outages.append({
                                                    'unit_name': unit_name,
                                                    'psr_type': unit_psr,
                                                    'psr_name': PSR_TYPES.get(unit_psr, unit_psr),
                                                    'capacity_mw': capacity_mw,
                                                    'start_time': pd.Timestamp(start_elem.text),
                                                    'end_time': pd.Timestamp(end_elem.text),
                                                    'businesstype': business_type,
                                                    'zone': zone
                                                })

                self._rate_limit()

            except Exception as e:
                # Empty response is OK (no outages)
                if "empty" not in str(e).lower():
                    print(f"    Warning: {zone} {psr_name} {start_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_outages:
            return pl.DataFrame(all_outages)
        else:
            return pl.DataFrame()

    def collect_generation_by_psr_type(
        self,
        zone: str,
        psr_type: str,
        start_date: str,
        end_date: str
    ) -> pl.DataFrame:
        """Collect generation for a specific PSR type.

        Args:
            zone: Bidding zone code
            psr_type: PSR type code (e.g., 'B04' for Gas, 'B14' for Nuclear)
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)

        Returns:
            Polars DataFrame with generation data for the PSR type
        """
        chunks = self._generate_monthly_chunks(start_date, end_date)
        all_data = []

        psr_name = PSR_TYPES.get(psr_type, psr_type)

        for start_chunk, end_chunk in tqdm(chunks, desc=f"  {zone} {psr_name}", leave=False):
            try:
                series = self.client.query_generation(
                    zone,
                    start=start_chunk,
                    end=end_chunk,
                    psr_type=psr_type
                )

                if series is not None and not series.empty:
                    # Handle both Series and DataFrame returns
                    if isinstance(series, pd.DataFrame):
                        series = series.iloc[:, 0]

                    df = pd.DataFrame({
                        'timestamp': series.index,
                        'generation_mw': series.values,
                        'zone': zone,
                        'psr_type': psr_type,
                        'psr_name': psr_name
                    })

                    pl_df = pl.from_pandas(df)
                    all_data.append(pl_df)

                self._rate_limit()

            except Exception as e:
                print(f"    Warning: {zone} {psr_name} {start_chunk.date()}: {e}")
                self._rate_limit()
                continue

        if all_data:
            return pl.concat(all_data)
        else:
            return pl.DataFrame()

    def collect_all(
        self,
        start_date: str,
        end_date: str,
        output_dir: Path
    ) -> dict:
        """Collect all ENTSO-E data with rate limiting.

        Args:
            start_date: Start date (YYYY-MM-DD)
            end_date: End date (YYYY-MM-DD)
            output_dir: Directory to save Parquet files

        Returns:
            Dictionary with paths to saved files
        """
        output_dir.mkdir(parents=True, exist_ok=True)

        # Calculate total requests
        months = len(self._generate_monthly_chunks(start_date, end_date))
        total_requests = (
            len(BIDDING_ZONES) * months * 2 +  # Generation + load
            len(BORDERS) * months  # Flows
        )
        estimated_minutes = total_requests / self.requests_per_minute

        print("=" * 70)
        print("ENTSO-E Data Collection")
        print("=" * 70)
        print(f"Date range: {start_date} to {end_date}")
        print(f"Bidding zones: {len(BIDDING_ZONES)}")
        print(f"Cross-border flows: {len(BORDERS)}")
        print(f"Monthly chunks: {months}")
        print(f"Total requests: ~{total_requests}")
        print(f"Rate limit: {self.requests_per_minute} requests/minute (45% of 60 max)")
        print(f"Estimated time: {estimated_minutes:.1f} minutes")
        print()

        results = {}

        # 1. Collect Generation Data
        print("[1/3] Collecting generation data by production type...")
        generation_data = []
        for zone in tqdm(BIDDING_ZONES.keys(), desc="Generation"):
            df = self.collect_generation_per_type(zone, start_date, end_date)
            if not df.is_empty():
                generation_data.append(df)

        if generation_data:
            generation_df = pl.concat(generation_data)
            gen_path = output_dir / "entsoe_generation_2024_2025.parquet"
            generation_df.write_parquet(gen_path)
            results['generation'] = gen_path
            print(f"✅ Generation: {generation_df.shape[0]:,} records → {gen_path}")

        # 2. Collect Load Data
        print("\n[2/3] Collecting load (demand) data...")
        load_data = []
        for zone in tqdm(BIDDING_ZONES.keys(), desc="Load"):
            df = self.collect_load(zone, start_date, end_date)
            if not df.is_empty():
                load_data.append(df)

        if load_data:
            load_df = pl.concat(load_data)
            load_path = output_dir / "entsoe_load_2024_2025.parquet"
            load_df.write_parquet(load_path)
            results['load'] = load_path
            print(f"✅ Load: {load_df.shape[0]:,} records → {load_path}")

        # 3. Collect Cross-Border Flows
        print("\n[3/3] Collecting cross-border flows...")
        flow_data = []
        for from_zone, to_zone in tqdm(BORDERS, desc="Flows"):
            df = self.collect_cross_border_flows(from_zone, to_zone, start_date, end_date)
            if not df.is_empty():
                flow_data.append(df)

        if flow_data:
            flow_df = pl.concat(flow_data)
            flow_path = output_dir / "entsoe_flows_2024_2025.parquet"
            flow_df.write_parquet(flow_path)
            results['flows'] = flow_path
            print(f"✅ Flows: {flow_df.shape[0]:,} records → {flow_path}")

        print()
        print("=" * 70)
        print("ENTSO-E Collection Complete")
        print("=" * 70)
        print(f"Total API requests made: {self.request_count}")
        print(f"Files created: {len(results)}")
        for data_type, path in results.items():
            file_size = path.stat().st_size / (1024**2)
            print(f"  - {data_type}: {file_size:.1f} MB")

        return results


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Collect ENTSO-E data with proper rate limiting")
    parser.add_argument(
        '--start-date',
        default='2024-10-01',
        help='Start date (YYYY-MM-DD)'
    )
    parser.add_argument(
        '--end-date',
        default='2025-09-30',
        help='End date (YYYY-MM-DD)'
    )
    parser.add_argument(
        '--output-dir',
        type=Path,
        default=Path('data/raw'),
        help='Output directory for Parquet files'
    )
    parser.add_argument(
        '--requests-per-minute',
        type=int,
        default=27,
        help='Requests per minute (default: 27 = 45%% of 60 limit)'
    )

    args = parser.parse_args()

    # Initialize collector and run
    collector = EntsoECollector(requests_per_minute=args.requests_per_minute)
    collector.collect_all(
        start_date=args.start_date,
        end_date=args.end_date,
        output_dir=args.output_dir
    )