# getfactormodels: https://github.com/x512/getfactormodels
# Copyright (C) 2025-2026 S. Martin <x512@pm.me>
# SPDX-License-Identifier: AGPL-3.0-or-later
#
# Distributed WITHOUT ANY WARRANTY. See LICENSE for full terms.
import io
import zipfile
from abc import ABC, abstractmethod
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.csv as pv
from typing_extensions import override
from getfactormodels.models.base import FactorModel, PortfolioBase, RegionMixin
from getfactormodels.utils.arrow_utils import (
    round_to_precision,
    scale_to_decimal,
)
from getfactormodels.utils.date_utils import offset_period_eom


#TODO: break up _read_zip
class FamaFrenchFactors(FactorModel, RegionMixin):
    """Download Fama-French (and Carhart) factor models.

    Downloads the 3-factor (1993), 5-factor (2015), or 6-factor (2018)
    model of Fama & French, or Carhart's (1997) 4-factor model.

    Weekly data only available for the 3-factor model.

    Args:
        model (str | int): model to return. '3' '4' '5' '6' (default: '3')
        frequency (str, optional): frequency of the data. 'd' 'm' 'y' 'w'
            (default: 'm').
        start_date (str, optional): start date of the data. YYYY[-MM-DD].
        end_date (str, optional): end date of the data. YYYY[-MM-DD].
        cache_ttl (int, optional`): cached download time-to-live in seconds 
            (default: 86400).
        region (str, optional): return an international/emerging market 
            model.

    References:
    - E. F. Fama and K. R., 1993. French. Common risk factors in the returns 
      on stocks and bonds. Journal of Financial Economics, vol. 33, no. 1,
      pp. 3–56.
    - E. F. Fama and K. R. French, 2015. A five-factor asset pricing model. 
      Journal of Financial Economics, vol. 116, no. 1, pp. 1–22.
    - E. F. Fama and K. R. French, 2018. Choosing factors. Journal of 
      Financial Economics, vol. 128, no. 2, pp. 234–252, 2018.

    """
    # TODO: NaNs in FamaFrench models! check
    #Note: "-0.9999" should be NaNs! [TODO: FIXME]
    @property
    def _frequencies(self) -> list[str]: return ['d', 'w', 'm', 'y']
    
    @property 
    def _precision(self) -> int: return 6

    @property 
    def _regions(self) -> list[str]:
        return [
            'us', 'emerging', 'developed', 'ex-us', 'europe',
            'japan', 'asia-pacific-ex-japan', 'north-america', 
        ]

    def __init__(self, 
                 frequency: str = 'm', 
                 model: int | str = '3', 
                 region: str | None = 'us', 
                 **kwargs) -> None:
        """Initialize a Fama-French factor model."""
        # model keys (ff3) to int
        self.model = model.replace('ff', '') if model.startswith('ff') else str(model)
        super().__init__(frequency=frequency, model=model, **kwargs)
        
        self.region = region
        self._validate_ff_input()

    @property
    def schema(self) -> pa.Schema:
        """Fama-French schema for specific model/frequency/region."""
        cols = [("date", pa.string()), 
                ("Mkt-RF", pa.float64()), 
                ("SMB", pa.float64()), 
                ("HML", pa.float64())]
        
        # 5 and 6-factor models: RMW and CMA
        if self.model in ["5", "6"]:
            cols += [("RMW", pa.float64()), 
                     ("CMA", pa.float64())]

        # mom models (4/6): dev/emerging: WML. US 6 factor: UMD. US 4: MOM.
        if self.model in ["4", "6"]:
            is_intl = self.region not in ["us", None]
            mom_name = "WML" if is_intl else ("UMD" if self.model == "6" else "MOM")
            cols.append((mom_name, pa.float64()))
        cols.append(("RF", pa.float64()))
        
        return pa.schema(cols)

    @property
    def _mom_schema(self) -> pa.Schema:
        """Private helper: schema for momentum files. 'val' is placeholder."""
        return pa.schema([("date", pa.string()), 
                          ("val", pa.float64())])
    @property
    def model(self) -> str:
        return self._model
    @model.setter
    def model(self, value: int | str):
        val = str(value)
        if val not in ["3", "4", "5", "6"]:
            raise ValueError(f"Invalid model '{val}'. Options: 3, 4, 5, 6")
        
        if hasattr(self, '_model') and val != self._model:
            self.log.info(f"Model changed from {self._model} to {val}")
            self._data = None
            self._view = None
        self._model = val 

    @property
    def _ff_region_map(self) -> dict[str, str]:
        """Private: maps region input to region URL str"""
        return {
            'us': 'US',
            'emerging': 'Emerging',
            'developed': 'Developed',
            'ex-us': 'Developed_ex_US',
            'europe': 'Europe',
            'japan': 'Japan',
            'asia-pacific-ex-japan': 'Asia_Pacific_ex_Japan',
            'north-america': 'North_America',
        }

    def _validate_ff_input(self) -> None:
        """Validates input for Fama-French models for things getters/setters dont take care of."""
        if self.region != 'us' and self.frequency == 'w':
            raise ValueError(f"Weekly frequency not available for {self.region}.")

        if self.region == "emerging" and self.frequency != 'm':
            raise ValueError("Emerging markets only available in monthly.")

        if self.frequency == 'w' and self.model != "3":
            raise ValueError("Weekly Fama-French data is only available for the 3-factor model.")
 

    def _get_url(self) -> dict[str, str]:
        """Constructs the URLs for downloading Fama-French data."""
        base_url = "https://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp"
        url_list = {}

        ff_url_name = self._ff_region_map.get(self.region)
        
        if ff_url_name == 'Emerging':
            # Emerging: only monthly 5-factors avail
            filename = "Emerging_5_Factors_CSV.zip"
        elif ff_url_name == 'US':
            _model = "F-F_Research_Data_5_Factors_2x3" if self.model in {"5", "6"} else "F-F_Research_Data_Factors"
            freq_map = {'d': '_daily', 'w': '_weekly'}
            suffix = freq_map.get(self.frequency, "")
            filename = f"{_model}{suffix}_CSV.zip"
        else:
            # International regions
            base_model = '3' if self.model in ['3', '4'] else '5'
            daily = '_Daily' if self.frequency == 'd' else ''
            filename = f"{ff_url_name}_{base_model}_Factors{daily}_CSV.zip"
            
        url_list['factors'] = f"{base_url}/{filename}"
        
        # Construct the momentum URL if needed.
        if self.model in ["4", "6"]:
            freq_suffix = "_daily" if self.frequency == 'd' else ""
            if self.region in ['us', None]:
                mom_filename = f"F-F_Momentum_Factor{freq_suffix}_CSV.zip"
            else:
                mom_filename = f"{self.region}_Mom_Factor{freq_suffix.title()}_CSV.zip"
            
            url_list['momentum'] = f"{base_url}/{mom_filename}"

        return url_list


    def _read_zip(self, _data: bytes, use_schema: pa.Schema = None) -> pa.Table:
        """Download and read CSV from zip file."""
        _schema = use_schema if use_schema else self.schema

        with zipfile.ZipFile(io.BytesIO(_data)) as z:
            filename = z.namelist()[0]
            with z.open(filename) as f:
                lines = f.read().decode('utf-8').splitlines()

        is_mom = 'momentum' in filename.lower() or '_mom_' in filename.lower()
        annual_marker = next((i for i, l in enumerate(lines) if "annual factors:" in l.lower()), None)

        if self.frequency == 'y':
            raw_content = lines[annual_marker + 1:] if annual_marker is not None else lines
        else:
            _start = 6 if self.region == 'emerging' else (12 if is_mom else 4)
            _stop = annual_marker if annual_marker else len(lines)
            raw_content = lines[_start:_stop]

        content = [",".join(_schema.names)]
        for line in raw_content:
            clean = line.strip()
            if not clean or "copyright" in clean.lower():
                if "copyright" in clean.lower(): self.copyright = clean
                continue
            
            # Use commas for split, not spaces!
            parts = [p.strip() for p in clean.split(',')]
            
            # only keeps lines that start with a date
            if parts and parts[0].isdigit():
                content.append(",".join(parts))

        return pv.read_csv(
            io.BytesIO("\n".join(content).encode('utf-8')),
            convert_options=pv.ConvertOptions(
                null_values=["-99.99", "-999", "-99.990", "-0.9999"],
                strings_can_be_null=True,
                column_types=_schema,
                include_columns=_schema.names,
            ),
        )   
    

    def _read(self, data: dict[str, bytes]) -> pa.Table:
        main_cols = [f for f in self.schema if f.name not in ["UMD", "MOM", "WML"]]
        
        # Emerging Markets: base file is 5 factors (even for models 3/4)
        if self.region == 'emerging' and self.model in ["3", "4"]:
            main_cols.insert(4, pa.field("RMW", pa.float64()))
            main_cols.insert(5, pa.field("CMA", pa.float64()))

        main_schema = pa.schema(main_cols)
        # data['factors']
        table = self._read_zip(data['factors'], use_schema=main_schema)

        if 'momentum' in data:
            mom_table = self._read_zip(data['momentum'], use_schema=self._mom_schema)
            
            # momentum column name based on schema
            mom_key = next(k for k in ["UMD", "MOM", "WML"] if k in self.schema.names)
            mom_table = mom_table.rename_columns(["date", mom_key])
            
            table = table.join(mom_table, keys="date", join_type="inner").combine_chunks()

        table = table.set_column(0, "date", table.column(0).cast(pa.string()))
        table = offset_period_eom(table, self.frequency)
        table = scale_to_decimal(table)
        
        table = round_to_precision(table, self._precision)

        return table.select(self.schema.names).combine_chunks()


# WIP ---------------------------------------------------------------------- #
# TODO: FactorModel to BaseModel.
class _FFPortfolioBase(PortfolioBase, ABC):
    """Base class for Fama-French portfolio return data."""
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    @abstractmethod
    def _find_table_start(self) -> str:
        """A pattern to identify the specific table's start."""
        ...

    def _get_ff_table(self, lines: list[str]) -> str:
        weight_ln = "Value" if self.weights == 'vw' else "Equal"
    
        # fix: US bivariate size/OP 2x3 sorts contain an error. Header line in the daily 
        # file is labelled 'Monthly' 
        is_size_op_2x3 = (
            hasattr(self, 'formed_on') and 
            frozenset(self.formed_on) == frozenset(['size', 'op']) and
            getattr(self, 'n_portfolios', None) in [6, "6"]
        )
        
        if is_size_op_2x3 and self.frequency == 'd':
            freq_label = "Monthly" 
        else:
            freq_label = "Annual" if self.frequency == 'y' else ("Daily" if self.frequency == 'd' else "Monthly")
        table_start = None
        for i, ln in enumerate(lines):
            if weight_ln in ln and "Returns" in ln and freq_label in ln:
                table_start = i
                break

        if table_start is None:
            raise ValueError(
                f"Could not find table header: `{weight_ln}`"
            )

        # Skip title line, find first non empty line
        header_idx = table_start + 1
        while header_idx < len(lines) and not lines[header_idx].strip():
            header_idx += 1

        if header_idx >= len(lines):
            raise ValueError("Couldn't find table header (found table title).")

        headers = ("date " + " ".join(lines[header_idx].split())).lower().replace("-", "_")
        cleaned = [headers]

        # Extract lines
        for line in lines[header_idx + 1:]:
            clean = line.strip()
            # Table's over if line doesn't start with a number
            if not clean or not clean[0].isdigit():
                break
            cleaned.append("  ".join(clean.split()))

        return "\n".join(cleaned)


    def _read(self, data: bytes) -> pa.Table:
        with zipfile.ZipFile(io.BytesIO(data)) as z:
            with z.open(z.namelist()[0]) as f:
                lines = f.read().decode('utf-8').splitlines()

        clean_csv = self._get_ff_table(lines)

        table = pv.read_csv(
            io.BytesIO(clean_csv.encode('utf-8')),
            convert_options=pv.ConvertOptions(
                null_values=["-99.99", "-999", "-99.9"],
                column_types={"date": pa.string()},
            ),
        )

        table = self._fix_ff_nulls(table)
        table = offset_period_eom(table, self.frequency)
        table = scale_to_decimal(table)

        renames = [n.strip().lower().replace(" ", "_").replace("-", "_") for n in table.column_names]
        table = table.rename_columns(renames)

        return table.combine_chunks()


    def _fix_ff_nulls(self, table: pa.Table) -> pa.Table:
        ff_nans = pa.array([-99.99, -999.0, -99.9], type=pa.float64())
        for i, name in enumerate(table.column_names):
            if name.lower() == "date": continue
            col = table.column(name)
            mask = pc.is_in(col, value_set=ff_nans)
            table = table.set_column(i, name, pc.if_else(mask, pa.scalar(None, type=col.type), col))
        return table


    @property
    def schema(self) -> pa.Schema:
        """Fama-French portfolio schema."""
        if self._data is not None:
            return self._data.schema
        return pa.schema([pa.field("date", pa.date32())])
    
    @property
    def _precision(self) -> int: return 6


class _FamaFrenchSorts(_FFPortfolioBase):
    @property
    def _frequencies(self) -> list[str]:
        return ['d', 'w', 'm', 'y']
    
    #@property 
    #def _regions(self) -> list[str]:
    #    return [
    #        'us', 'emerging', 'developed', 'ex-us', 'europe',
    #        'japan', 'asia-pacific-ex-japan', 'north-america', 
    #    ]
    # TODO: regions
    def __init__(self,
                 formed_on: str | list[str] = 'size',
                 sort: str | int | None = None,
    #            region: str = 'us', 
                 **kwargs):
        super().__init__(**kwargs)
        
        #self.region = region

        # formed_on: sorted list of lowercase strings
        if isinstance(formed_on, str):
            self.formed_on = [formed_on.lower()]
        else:
            self.formed_on = sorted([f.lower() for f in formed_on])

        self.sort = str(sort).lower() if sort is not None else self._get_default_sort()
        self.n_portfolios = self._map_sort_to_count()
        self.is_multivariate = len(self.formed_on) > 1

        self._validate_params()


    def _find_table_start(self) -> str:
        """Identify the table section based on weights."""
        return "Equal Weighted Returns" if self.weights == 'ew' else "Value Weighted Returns"


    def _map_sort_to_count(self) -> int:
        """Map sort string/int to total portfolio count."""
        mapping = {
            'tertile': 3, '3': 3,
            'quintile': 5, '5': 5,
            'decile': 10, '10': 10,
            '2x3': 6, '6': 6,
            '5x5': 25, '25': 25,
            '10x10': 100, '100': 100,
            '2x4x4': 32, '32': 32,
        }
        if self.sort not in mapping:
            raise ValueError(f"Sort '{self.sort}' not recognized. Try 'decile', '5x5', etc.")
        return mapping[self.sort]


    def _get_default_sort(self) -> str:
        """Assigns a default sort if none provided."""
        if len(self.formed_on) == 1:
            return 'decile'
        if len(self.formed_on) == 2:
            if any(f in {'ep', 'cfp', 'dp'} for f in self.formed_on):
                return '2x3'
            return '5x5'
        return '2x4x4'


    def _validate_params(self):
        """Validates factors, frequency and sort types."""
        valid = {'size', 'bm', 'op', 'inv', 'ep', 'cfp', 'dp', 'mom',
                 'st_rev', 'lt_rev', 'ac', 'beta', 'ni', 'var', 'resvar'}
        for f in self.formed_on:
            if f not in valid:
                raise ValueError(f"Invalid factor: {f}. Must be one of {sorted(valid)}")

        no_daily = {'ep', 'cfp', 'dp', 'beta', 'ac', 'ni', 'var', 'resvar'}
        if self.frequency == 'd' and any(f in no_daily for f in self.formed_on):
            raise ValueError(f"Daily data is not available for sorts involving {self.formed_on}")

        prior_rets = {'mom', 'st_rev', 'lt_rev'}
        if not self.is_multivariate:
            # Univariate sorts only support 3, 5, 10
            # This fixes eg, "-p 2x3 --on op" returning the full table.
            if self.n_portfolios not in {3, 5, 10}:
                 raise ValueError(
                     f"Sort '{self.sort}' ({self.n_portfolios}) is not a valid univariate sort. "
                     f"Use 'tertile', 'quintile', 'decile' (or 3, 5 or 10)."
                 )

            if self.formed_on[0] in prior_rets: #noqa
                if self.n_portfolios != 10:
                    raise ValueError(f"{self.formed_on[0]} is only available in deciles.")
    
    @override
    def _read(self, data: bytes) -> pa.Table:
        table = super()._read(data)
        
        # fix: drop the "<= 0" column here (bivariates can contain it)
        neg_cols = [name for name in table.column_names if "<=" in name]
        if neg_cols:
            table = table.drop(neg_cols) #

        # Sorts on prior rets are decile only, the rest are a table 
        # of date [+ <= 0] + tertiles, quintiles, deciles.
        if not self.is_multivariate and self.formed_on[0] not in {'mom', 'st_rev', 'lt_rev'}:
            table = self._slice_univariate(table)
        
        return table.combine_chunks()


    def _slice_univariate(self, table: pa.Table) -> pa.Table:
        # No "< =" col now
        mapping = {
            3: slice(1, 4),
            5: slice(4, 9),
            10: slice(9, 19),
        }
        
        target_slice = mapping.get(self.n_portfolios)

        if target_slice is None:
             raise ValueError(f"Cannot slice univariate table: no mapping for {self.n_portfolios} portfolios.")

        if target_slice:
            if table.num_columns > 16:
                indices = [0] + list(range(target_slice.start, target_slice.stop))
                return table.select([i for i in indices if i < table.num_columns])
            # quick fix: 'ac' and 'beta' univariates don't have tertiles... 
            if table.num_columns >= 15:
                adj_mapping = {
                    5: slice(table.num_columns - 15, table.num_columns - 10),
                    10: slice(table.num_columns - 10, table.num_columns),
                }
                adj_slice = adj_mapping.get(self.n_portfolios)
                if adj_slice:
                    indices = [0] + list(range(adj_slice.start, adj_slice.stop))
                    return table.select([i for i in indices if i < table.num_columns])

        return table


    def _get_url(self) -> str:
        """Constructs a URL for a portfolio .zip on the Ken French library."""
        base = "https://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp"

        if self.frequency == 'd':
            suffix = "_Daily_CSV.zip"
        elif not self.dividends:
            suffix = "_Wout_Div_CSV.zip"
        else:
            suffix = "_CSV.zip"

        # Univariate sorts
        if not self.is_multivariate:
            mapping = {
                'size': 'Portfolios_Formed_on_ME',
                'bm': 'Portfolios_Formed_on_BE-ME',
                'op': 'Portfolios_Formed_on_OP',
                'inv': 'Portfolios_Formed_on_INV',
                'ep': 'Portfolios_Formed_on_E-P',
                'cfp': 'Portfolios_Formed_on_CF-P',
                'dp': 'Portfolios_Formed_on_D-P',
                'mom': '10_Portfolios_Prior_12_2',
                'st_rev': '10_Portfolios_Prior_1_0',
                'lt_rev': '10_Portfolios_Prior_60_13',
                'ac': 'Portfolios_Formed_on_AC',
                'beta': 'Portfolios_Formed_on_BETA',
                'ni': 'Portfolios_Formed_on_NI',
                'var': 'Portfolios_Formed_on_VAR',
                'resvar': 'Portfolios_Formed_on_RESVAR',
            }
            return f"{base}/{mapping[self.formed_on[0]]}{suffix}"

        # Multivariate sorts
        factors = frozenset(self.formed_on)

        # 3-way sorts  # NOTE REGIONAL AVAIL.
        if len(factors) == 3:
            three_way = {
                frozenset(['size', 'bm', 'op']): "32_Portfolios_ME_BEME_OP_2x4x4",
                frozenset(['size', 'bm', 'inv']): "32_Portfolios_ME_BEME_INV_2x4x4",
                frozenset(['size', 'op', 'inv']): "32_Portfolios_ME_OP_INV_2x4x4",
            }
            slug = three_way.get(factors)
            if not slug: 
                raise ValueError("3-way sort not supported.")

            return f"{base}/{slug}{suffix}"

        # Bivariate sorts
        bivariate_map = {
            frozenset(['size', 'bm']): "",
            frozenset(['size', 'op']): "ME_OP",
            frozenset(['size', 'inv']): "ME_INV",
            frozenset(['bm', 'op']): "BEME_OP",
            frozenset(['bm', 'inv']): "BEME_INV",
            frozenset(['op', 'inv']): "OP_INV",
            frozenset(['size', 'ep']): "ME_EP",
            frozenset(['size', 'cfp']): "ME_CFP",
            frozenset(['size', 'dp']): "ME_DP",
            frozenset(['size', 'ac']): "ME_AC",
            frozenset(['size', 'beta']): "ME_BETA",
            frozenset(['size', 'ni']): "ME_NI",
            frozenset(['size', 'var']): "ME_VAR",
            frozenset(['size', 'resvar']): "ME_RESVAR",
            frozenset(['size', 'mom']): "ME_Prior_12_2",
            frozenset(['size', 'st_rev']): "ME_Prior_1_0",
            frozenset(['size', 'lt_rev']): "ME_Prior_60_13",
        }

        slug_part = bivariate_map.get(factors)
        if slug_part is None:
            raise ValueError(f"Bivariate combination {self.formed_on} not supported.")

        grid_map = {6: ("6_Portfolios", "2x3"),
                    25: ("25_Portfolios", "5x5"),
                    100: ("100_Portfolios", "10x10")}
        prefix, grid = grid_map.get(self.n_portfolios, ("25_Portfolios", "5x5"))

        if not slug_part: # bivariate
            return f"{base}/{prefix}_{grid}{suffix}"
        if "Prior" in slug_part: # bivariate sorts on prior returns
            return f"{base}/{prefix}_{slug_part}{suffix}"

        return f"{base}/{prefix}_{slug_part}_{grid}{suffix}"


class _FamaFrenchIndustryPortfolios(_FFPortfolioBase):
    """Fama-French industry portfolios."""
    _SIZES = {5, 10, 12, 17, 30, 38, 48, 49}

    @property
    def _frequencies(self) -> list[str]: return ['d', 'm', 'y']

    def __init__(self, n_portfolios: int | str = 5, **kwargs):
        n_portfolios = int(n_portfolios) # allow str through
        if n_portfolios not in self._SIZES:
            raise ValueError(f"Industry count must be one of {self._SIZES}")
        self.n_portfolios = n_portfolios
        super().__init__(**kwargs)


    def _find_table_start(self) -> str:
        w_str = "Value Weighted" if self.weights == 'vw' else "Equal Weighted"
        freq_map = {'y': 'Annual', 'd': 'Daily', 'm': 'Monthly'}

        return f"Average {w_str} Returns -- {freq_map.get(self.frequency, 'Monthly')}"


    def _get_url(self) -> str:
        base = "https://mba.tuck.dartmouth.edu/pages/faculty/ken.french/ftp"

        if self.frequency == 'd':
            suffix = "daily_CSV.zip"
        else:
            suffix = "Wout_Div_CSV.zip" if not self.dividends else "CSV.zip"

        return f"{base}/{self.n_portfolios}_Industry_Portfolios_{suffix}"


def _get_ff_portfolios(
    formed_on: str | list[str] = 'size', #on=, sorted_on=, might rename
    sort: str | int | None = None, #2x3, decile, 25
    industry: int | None = None, 
    weights: str = 'vw',
    frequency: str = 'm',
    **kwargs) -> pa.Table:
    """Internal helper: factory to return.

    Factory for Fama French portfolios. Either `industry` or 
    `formed_on` required.
    """
    # we industry?
    if industry is not None:
        return _FamaFrenchIndustryPortfolios(
            frequency=frequency, weights=weights, n_portfolios=industry, **kwargs
        )

    # keep formed_on None, let class handle it.
    return _FamaFrenchSorts(
        frequency=frequency, weights=weights, formed_on=formed_on, sort=sort, **kwargs
    )

# TODO: setters/getters, regions, tests, aggregates for q? prob not. 
