Source code for kaxanuk.attribution_analysis.input_handlers.csv_portfolio_input_handler

"""
CSV portfolio input handler for reading portfolio allocation files.

Read portfolio weight files in CSV format, supporting both horizontal
(Ticker | Date1 | Date2 | ...) and vertical (Date | Ticker1 | Ticker2 | ...)
layouts with automatic format detection and transposition.
"""

from pathlib import Path
import csv as python_csv
import logging

import pyarrow as pa
from pyarrow import csv
import pyarrow.compute as pc

from kaxanuk.attribution_analysis.exceptions import MissingPortfolioError, ReadingCsvPortfolioError
from kaxanuk.attribution_analysis.input_handlers import input_handlers_helpers
from kaxanuk.attribution_analysis.portfolio_data_pipeline import portfolio_transformer
from kaxanuk.attribution_analysis.input_handlers.portfolio_input_handler_interface import PortfolioInputHandlerInterface

MAX_PRECISION = 38
PRECISION = 72

logger = logging.getLogger(__name__)



class CsvPortfolioInputHandler(PortfolioInputHandlerInterface):
    """
    Handles reading and parsing of CSV portfolio files.

    This class reads a portfolio CSV file, infers decimal precision
    dynamically for numeric columns, and converts them into a typed
    PyArrow table suitable for financial pipelines.

    Supports both horizontal and vertical formats:
        - Horizontal: Ticker | Date1 | Date2 | ...
        - Vertical: Date | Ticker1 | Ticker2 | ... (auto-transposes to horizontal)

    Attributes
    ----------
    base_dir : Path
        Directory where portfolio CSV files are located.
    """

    def __init__(
            self,
            base_dir: str
    ) -> None:
        """
        Initialize the CSV portfolio handler.

        Parameters
        ----------
        base_dir : str
            Base directory containing portfolio CSV files.
        """
        self.base_dir = Path(base_dir)



    def read_portfolio(
            self,
            *,
            file_name: str
    ) -> pa.Table:
        """
        Reads a portfolio from a CSV file and converts it into a typed PyArrow Table.

        Automatically detects horizontal vs vertical format and returns horizontal format.

        Parameters
        ----------
        file_name : str
            Name of the CSV file (with or without .csv extension).

        Returns
        -------
        pa.Table
            Processed PyArrow table containing portfolio tickers and weights in horizontal format.

        Raises
        ------
        FileNotFoundError
            If the file does not exist in the specified base directory.

        ReadingCsvPortfolioError
            If the CSV cannot be read or processed correctly.
        """
        file_path = (
            self.base_dir / file_name
            if file_name.lower().endswith(".csv")
            else self.base_dir / f"{file_name}.csv"
        )

        if not file_path.exists():
            raise MissingPortfolioError(
                f"Portfolio file not found in {self.base_dir}: {file_path.name}"
            )

        new_columns: list[pa.ChunkedArray] = []
        csv_column_names: list[str] = []

        try:
            # Detect format
            with file_path.open() as f:
                reader = python_csv.reader(f)
                headers = next(reader)

            first_col = headers[0].strip()
            format_type = portfolio_transformer.detect_portfolio_format(
                first_column_name=first_col,
                error_class=ReadingCsvPortfolioError
            )
            logger.info("Detected %s format for %s", format_type, file_path.name)

            if format_type == "vertical":
                with file_path.open() as f:
                    reader = python_csv.reader(f)
                    headers = next(reader)

                # Force first column (Date) to be string, rest can be anything
                column_types = {headers[0]: pa.string()}
                for header in headers[1:]:
                    column_types[header] = pa.string()

                convert_options = csv.ConvertOptions(column_types=column_types)
                prem_table = csv.read_csv(file_path, convert_options=convert_options)

                # Transpose to horizontal format
                prem_table = portfolio_transformer.transpose_vertical_to_horizontal(table=prem_table)
            else:
                # Horizontal format - read with all string types
                convert_options = csv.ConvertOptions(
                    column_types={h.strip(): pa.string() for h in headers}
                )
                prem_table = csv.read_csv(file_path, convert_options=convert_options)

                # Normalize first column to "Ticker" (case-insensitive)
                if prem_table.column_names[0].lower() == "ticker":
                    new_names = ["Ticker", *prem_table.column_names[1:]]
                    prem_table = prem_table.rename_columns(new_names)

            # Guard: first column must be "Ticker" after normalization
            if prem_table.column_names[0] != "Ticker":
                raise ReadingCsvPortfolioError(
                    f"First column of portfolio must be named 'Ticker' "
                    f"(found: '{prem_table.column_names[0]}')"
                )

            for column_name in prem_table.column_names:
                csv_column_names.append(column_name)
                col = prem_table[column_name]

                if column_name == "Ticker":
                    new_columns.append(col)
                else:
                    decimal_spec = input_handlers_helpers.infer_decimal_precision_and_scale(col=col)
                    target_type = input_handlers_helpers.get_decimal_type(
                        precision=decimal_spec.precision,
                        scale=decimal_spec.scale,
                        precision_limit=PRECISION
                    )
                    new_columns.append(pc.cast(col, target_type, safe=False))
            return pa.table(new_columns, names=csv_column_names)

        except ReadingCsvPortfolioError:
            raise
        except FileNotFoundError as error:
            raise ReadingCsvPortfolioError(
                f"Portfolio file not found: {file_path}"
            ) from error
        except pa.lib.ArrowInvalid as error:
            raise ReadingCsvPortfolioError(
                f"Invalid CSV format in portfolio file {file_path.name}: {error}"
            ) from error
        except (ValueError, KeyError, TypeError) as error:
            raise ReadingCsvPortfolioError(
                f"Error parsing portfolio file {file_path.name}: {error}"
            ) from error
        except OSError as error:
            raise ReadingCsvPortfolioError(
                f"I/O error reading portfolio file {file_path.name}: {error}"
            ) from error