<commit_before>"""
Aggregate plant parts to make an EIA master plant-part table.

Practically speaking, a plant is a collection of generator(s). There are many
attributes of generators (i.e. prime mover, primary fuel source, technology
type). We can use these generator attributes to group generator records into
larger aggregate records which we call "plant-parts". A plant part is a record
which corresponds to a particular collection of generators that all share an
identical attribute. E.g. all of the generators with unit_id=2, or all of the
generators with coal as their primary fuel source.

The EIA data about power plants (from EIA 923 and 860) is reported in tables
with records that correspond to mostly generators and plants. Other datasets
(cough cough FERC1) are less well organized and include plants, generators and
other plant-parts all in the same table without any clear labels. The master
plant-part table is an attempt to create records corresponding to many
different plant-parts in order to connect specific slices of EIA plants to
other datasets.

Because generators are often owned by multiple utilities, another dimention of
the master unit list involves generating two records for each owner: one of the
portion of the plant part they own and one for the plant part as a whole. The
portion records are labeled in the ``ownership`` column as "owned" and the total
records are labeled as "total".

This module refers to "true granularies". Many plant parts we cobble together
here in the master plant-part list refer to the same collection of
infrastructure as other plant-part list records. For example, if we have a
"plant_prime_mover" plant part record and a "plant_unit" plant part record
which were both cobbled together from the same two generators. We want to be
able to reduce the plant-part list to only unique collections of generators,
so we label the first unique granularity as a true granularity and label the
subsequent records as false granularities with the ``true_gran`` column. In
order to choose which plant-part to keep in these instances, we assigned a
:py:const:`PLANT_PARTS_ORDERED` and label whichever plant-part comes first as
the unique granularity.

**Recipe Book for the plant-part list**

:py:const:`PLANT_PARTS` is the main recipe book for how each of the plant-parts
need to be compiled. These plant-parts represent ways to group generators based
on widely reported values in EIA. All of these are logical ways to group
collections of generators - in most cases - but some groupings of generators
are more prevelant or relevant than others for certain types of plants.

The canonical example here is the ``plant_unit``. A unit is a collection of
generators that operate together - most notably the combined-cycle natural gas
plants. Combined-cycle units generally consist of a number of gas turbines
which feed excess steam to a number of steam turbines.

>>> df_gens = pd.DataFrame({
...     'plant_id_eia': [1, 1, 1],
...     'generator_id': ['a', 'b', 'c'],
...     'unit_id_pudl': [1, 1, 1],
...     'prime_mover_code': ['CT', 'CT', 'CA'],
...     'capacity_mw': [50, 50, 100],
... })
>>> df_gens
    plant_id_eia    generator_id    unit_id_pudl    prime_mover_code    capacity_mw
0              1               a               1                  CT             50
1              1               b               1                  CT             50
2              1               c               1                  CA            100

A good example of a plant-part that isn't really logical also comes from a
combined-cycle unit. Grouping this example plant by the ``prime_mover_code``
would generate two records that would basically never show up in FERC1.
This stems from the inseparability of the generators.

>>> df_plant_prime_mover = pd.DataFrame({
...     'plant_id_eia': [1, 1],
...     'plant_part': ['plant_prime_mover', 'plant_prime_mover'],
...     'prime_mover_code': ['CT', 'CA'],
...     'capacity_mw': [100, 100],
... })
>>> df_plant_prime_mover
    plant_id_eia         plant_part    prime_mover_code    capacity_mw
0              1  plant_prime_mover                  CT            100
1              1  plant_prime_mover                  CA            100

In this case the unit is more relevant:

>>> df_plant_unit = pd.DataFrame({
...     'plant_id_eia': [1],
...     'plant_part': ['plant_unit'],
...     'unit_id_pudl': [1],
...     'capacity_mw': [200],
... })
>>> df_plant_unit
    plant_id_eia    plant_part    unit_id_pudl    capacity_mw
0              1    plant_unit               1            200

But if this same plant had both this combined-cycle unit and two more
generators that were self contained "GT" or gas combustion turbine, a frequent
way to group these generators is differnt for the combined-cycle unit and the
gas-turbine.

>>> df_gens = pd.DataFrame({
...     'plant_id_eia': [1, 1, 1, 1, 1],
...     'generator_id': ['a', 'b', 'c', 'd', 'e'],
...     'unit_id_pudl': [1, 1, 1, 2, 3],
...     'prime_mover_code': ['CT', 'CT', 'CA', 'GT', 'GT'],
...     'capacity_mw': [50, 50, 100, 75, 75],
... })
>>> df_gens
    plant_id_eia    generator_id    unit_id_pudl    prime_mover_code    capacity_mw
0              1               a               1                  CT             50
1              1               b               1                  CT             50
2              1               c               1                  CA            100
3              1               d               2                  GT             75
4              1               e               3                  GT             75

>>> df_plant_part = pd.DataFrame({
...     'plant_id_eia': [1, 1],
...     'plant_part': ['plant_unit', 'plant_prime_mover'],
...     'unit_id_pudl': [1, pd.NA],
...     'prime_mover_code': [pd.NA, 'GT',],
...     'capacity_mw': [200, 150],
... })
>>> df_plant_part
    plant_id_eia           plant_part    unit_id_pudl    prime_mover_code    capacity_mw
0              1           plant_unit               1                <NA>            200
1              1    plant_prime_mover            <NA>                  GT            150

In this case last, the ``plant_unit`` record would have a null
``plant_prime_mover`` because the unit contains more than one
``prime_mover_code``. Same goes for the ``unit_id_pudl`` of the
``plant_prime_mover``. This is handled in the :class:``AddConsistentAttributes``.

**Overview of flow for generating the master unit list:**

The three main classes which enable the generation of the plant-part table are:

* :class:`MakeMegaGenTbl`: All of the plant parts are compiled from generators.
  So this class generates a big dataframe of generators with any ID and data
  columns we'll need. This is also where we add records regarding utility
  ownership slices. The table includes two records for every generator-owner:
  one for the "total" generator (assuming the owner owns 100% of the generator)
  and one for the report ownership fraction of that generator with all of the
  data columns scaled to the ownership fraction.
* :class:`LabelTrueGranularities`: This class creates labels for all generators
  which note wether the plant-part records that will be compiled from each
  generator will be a "true granulary", as described above.
* :class:`MakePlantParts`: This class uses the generator dataframe, the
  granularity dataframe from the above two classes as well as the information
  stored in :py:const:`PLANT_PARTS` to know how to aggregate each of the plant
  parts. Then we have plant part dataframes with the columns which identify the
  plant part and all of the data columns aggregated to the level of the plant
  part. With that compiled plant part dataframe we also add in qualifier
  columns with :class:`AddConsistentAttributes`. A qualifer column is a column which
  contain data that is not endemic to the plant part record (it is not one of
  the identifying columns or aggregated data columns) but the data is still
  useful data that is attributable to each of the plant part records. For more
  detail on what a qualifier column is, see :meth:`AddConsistentAttributes.execute`.

**Generating the plant-parts list**

There are two ways to generate the plant-parts table: one directly using the
:class:`pudl.output.pudltabl.PudlTabl` object and the other using the classes
from this module. Either option needs a :class:`pudl.output.pudltabl.PudlTabl`
object.

Create the :class:`pudl.output.pudltabl.PudlTabl` object:

.. code-block:: python

    import pudl
    pudl_engine = sa.create_engine(pudl.workspace.setup.get_defaults()['pudl_db'])
    pudl_out = pudl.output.pudltabl.PudlTabl(pudl_engine,freq='AS')

Then make the table via pudl_out:

.. code-block:: python

    plant_parts_eia = pudl_out.plant_parts_eia()


OR make the table via objects in this module:

.. code-block:: python

    gens_mega = MakeMegaGenTbl().execute(mcoe, own_eia860)
    true_grans = LabelTrueGranularities().execute(gens_mega)
    parts_compiler = MakePlantParts(pudl_out)
    plant_parts_eia = parts_compiler.execute(gens_mega=gens_mega, true_grans=true_grans)


"""

import logging
import warnings
from copy import deepcopy
from typing import Dict, List, Literal

import numpy as np
import pandas as pd

import pudl

logger = logging.getLogger(__name__)

# HALP: I need both of these setting set in order for the dfs in the docstrings
# to pass the doctests. Without them the formatting get all jumbled.
# but obviously this is the wrong place to do this.
# I tried adding these into conftest.py in pandas_terminal_width().
# I tried adding this into __init__.py.
# I tried adding this into the module docstring.
pd.options.display.width = 1000
pd.options.display.max_columns = 1000

PLANT_PARTS: Dict[str, Dict[str, List]] = {
    "plant": {
        "id_cols": ["plant_id_eia"],
    },
    "plant_gen": {
        "id_cols": ["plant_id_eia", "generator_id"],
    },
    "plant_unit": {
        "id_cols": ["plant_id_eia", "unit_id_pudl"],
    },
    "plant_technology": {
        "id_cols": ["plant_id_eia", "technology_description"],
    },
    "plant_prime_fuel": {  # 'plant_primary_fuel': {
        "id_cols": ["plant_id_eia", "energy_source_code_1"],
    },
    "plant_prime_mover": {
        "id_cols": ["plant_id_eia", "prime_mover_code"],
    },
    "plant_ferc_acct": {
        "id_cols": ["plant_id_eia", "ferc_acct_name"],
    },
}
"""
dict: this dictionary contains a key for each of the 'plant parts' that should
end up in the mater unit list. The top-level value for each key is another
dictionary, which contains keys:

* id_cols (the primary key type id columns for this plant part). The
  plant_id_eia column must come first.

"""

PLANT_PARTS_ORDERED: List[str] = [
    "plant",
    "plant_unit",
    "plant_prime_mover",
    "plant_technology",
    "plant_prime_fuel",
    "plant_ferc_acct",
    "plant_gen",
]


IDX_TO_ADD: List[str] = ["report_date", "operational_status_pudl"]
"""
list: list of additional columns to add to the id_cols in :py:const:`PLANT_PARTS`.
The id_cols are the base columns that we need to aggregate on, but we also need
to add the report date to keep the records time sensitive and the
operational_status_pudl to separate the operating plant-parts from the
non-operating plant-parts.
"""

IDX_OWN_TO_ADD: List[str] = ["utility_id_eia", "ownership"]
"""
list: list of additional columns beyond the :py:const:`IDX_TO_ADD` to add to the
id_cols in :py:const:`PLANT_PARTS` when we are dealing with plant-part records
that have been broken out into "owned" and "total" records for each of their
owners.
"""

SUM_COLS: List[str] = [
    "total_fuel_cost",
    "net_generation_mwh",
    "capacity_mw",
    "capacity_eoy_mw",
    "total_mmbtu",
]
"""
list: list of columns to sum when aggregating a table.
"""

WTAVG_DICT = {
    "fuel_cost_per_mwh": "capacity_mw",
    "heat_rate_mmbtu_mwh": "capacity_mw",
    "fuel_cost_per_mmbtu": "capacity_mw",
}
"""
dict: a dictionary of columns (keys) to perform weighted averages on and
the weight column (values)"""


CONSISTENT_ATTRIBUTE_COLS = [
    "fuel_type_code_pudl",
    "planned_retirement_date",
    "retirement_date",
    "generator_id",
    "unit_id_pudl",
    "technology_description",
    "energy_source_code_1",
    "prime_mover_code",
    "ferc_acct_name",
]
"""
list: a list of column names to add as attributes when they are consistent into
the aggregated plant-part records.
"""

PRIORITY_ATTRIBUTES_DICT = {
    "operational_status": ["existing", "proposed", "retired"],
}


FIRST_COLS = [
    "plant_id_eia",
    "report_date",
    "plant_part",
    "generator_id",
    "unit_id_pudl",
    "prime_mover_code",
    "energy_source_code_1",
    "technology_description",
    "ferc_acct_name",
    "utility_id_eia",
    "true_gran",
    "appro_part_label",
]


class MakeMegaGenTbl(object):
    """
    Compiler for a MEGA generator table with ownership integrated.

    Examples
    --------
    **Input Tables**

    Here is an example of one plant with three generators. We will use
    ``capacity_mw`` as the data column.

    >>> mcoe = pd.DataFrame({
    ...     'plant_id_eia': [1, 1, 1],
    ...     'report_date': ['2020-01-01', '2020-01-01','2020-01-01'],
    ...     'generator_id': ['a', 'b', 'c'],
    ...     'utility_id_eia': [111, 111, 111],
    ...     'unit_id_pudl': [1, 1, 1],
    ...     'prime_mover_code': ['CT', 'CT', 'CA'],
    ...     'technology_description': [
    ...         'Natural Gas Fired Combined Cycle', 'Natural Gas Fired Combined Cycle', 'Natural Gas Fired Combined Cycle'
    ...     ],
    ...     'operational_status': ['existing', 'existing','existing'],
    ...     'retirement_date': [pd.NA, pd.NA, pd.NA],
    ...     'capacity_mw': [50, 50, 100],
    ... }).astype({
    ...     'retirement_date': "datetime64[ns]",
    ...     'report_date': "datetime64[ns]",
    ... })
    >>> mcoe
        plant_id_eia    report_date 	generator_id   utility_id_eia 	unit_id_pudl 	prime_mover_code              technology_description   operational_status  retirement_date 	capacity_mw
    0 	           1     2020-01-01 	           a              111 	           1 	              CT    Natural Gas Fired Combined Cycle             existing               NaT 	         50
    1 	           1     2020-01-01 	           b              111 	           1 	              CT    Natural Gas Fired Combined Cycle             existing               NaT 	         50
    2 	           1     2020-01-01 	           c              111 	           1 	              CA    Natural Gas Fired Combined Cycle             existing               NaT 	        100

    The ownership table from EIA 860 includes one record for every owner of
    each generator. In this example generator ``c`` has two owners.

    >>> df_own_eia860 = pd.DataFrame({
    ...     'plant_id_eia': [1, 1, 1, 1],
    ...     'report_date': ['2020-01-01', '2020-01-01','2020-01-01', '2020-01-01'],
    ...     'generator_id': ['a', 'b', 'c', 'c'],
    ...     'utility_id_eia': [111, 111, 111, 111],
    ...     'owner_utility_id_eia': [111, 111, 111, 888],
    ...     'fraction_owned': [1, 1, .75, .25]
    ... }).astype({'report_date': "datetime64[ns]"})
    >>> df_own_eia860
        plant_id_eia 	report_date   generator_id 	utility_id_eia 	owner_utility_id_eia  fraction_owned
    0 	           1 	 2020-01-01 	         a 	           111 	                 111 	        1.00
    1 	           1 	 2020-01-01 	         b 	           111 	                 111 	        1.00
    2 	           1 	 2020-01-01 	         c 	           111 	                 111 	        0.75
    3 	           1 	 2020-01-01 	         c 	           111 	                 888 	        0.25

    **Output Mega Generators Table**

    ``MakeMegaGenTbl().execute(mcoe, df_own_eia860, slice_cols=['capacity_mw'])``
    produces the output table ``gens_mega`` which includes two main sections:
    the generators with a "total" ownership stake for each of their owners and
    the generators with an "owned" ownership stake for each of their
    owners. For the generators that are owned 100% by one utility, the
    records are identical except the ``ownership`` column. For the
    generators that have more than one owner, there are two "total" records
    with 100% of the capacity of that generator - one for each owner - and
    two "owned" records with the capacity scaled to the ownership stake
    of each of the owner utilites - represented by ``fraction_owned``.

    """

    def __init__(self):
        """
        Initialize object which creates a MEGA generator table.

        The coordinating function here is :meth:`execute`.

        """
        self.id_cols_list = make_id_cols_list()

    def execute(
        self,
        mcoe: pd.DataFrame,
        own_eia860: pd.DataFrame,
        slice_cols: List[str] = SUM_COLS,
        validate_own_merge: Literal["1:1", "1:m", "m:m"] = "1:m",
    ) -> pd.DataFrame:
        """
        Make the mega generators table with ownership integrated.

        Args:
            mcoe: generator-based mcoe table from :meth:`pudl.output.PudlTabl.mcoe()`
            own_eia860: ownership table from :meth:`pudl.output.PudlTabl.own_eia860()`
            slice_cols: list of columns to slice by ownership fraction in
                :meth:`MakeMegaGenTbl.slice_by_ownership`. Default is :py:const:`SUM_COLS`

        Returns:
            a table of all of the generators with identifying columns and data
            columns, sliced by ownership which makes "total" and "owned"
            records for each generator owner. The "owned" records have the
            generator's data scaled to the ownership percentage (e.g. if a 200
            MW generator has a 75% stake owner and a 25% stake owner, this will
            result in two "owned" records with 150 MW and 50 MW). The "total"
            records correspond to the full plant for every owner (e.g. using
            the same 2-owner 200 MW generator as above, each owner will have a
            records with 200 MW).

        """
        logger.info("Generating the mega generator table with ownership.")

        gens_mega = (
            self.get_gens_mega_table(mcoe)
            .pipe(self.label_operating_gens)
            .pipe(self.slice_by_ownership, own_eia860, slice_cols, validate_own_merge)
        )
        return gens_mega

    def get_gens_mega_table(self, mcoe):
        """
        Compile the main generators table that will be used as base of PPL.<commit_msg>update validation arg in ppe merge<commit_after>"""
Aggregate plant parts to make an EIA master plant-part table.

Practically speaking, a plant is a collection of generator(s). There are many
attributes of generators (i.e. prime mover, primary fuel source, technology
type). We can use these generator attributes to group generator records into
larger aggregate records which we call "plant-parts". A plant part is a record
which corresponds to a particular collection of generators that all share an
identical attribute. E.g. all of the generators with unit_id=2, or all of the
generators with coal as their primary fuel source.

The EIA data about power plants (from EIA 923 and 860) is reported in tables
with records that correspond to mostly generators and plants. Other datasets
(cough cough FERC1) are less well organized and include plants, generators and
other plant-parts all in the same table without any clear labels. The master
plant-part table is an attempt to create records corresponding to many
different plant-parts in order to connect specific slices of EIA plants to
other datasets.

Because generators are often owned by multiple utilities, another dimention of
the master unit list involves generating two records for each owner: one of the
portion of the plant part they own and one for the plant part as a whole. The
portion records are labeled in the ``ownership`` column as "owned" and the total
records are labeled as "total".

This module refers to "true granularies". Many plant parts we cobble together
here in the master plant-part list refer to the same collection of
infrastructure as other plant-part list records. For example, if we have a
"plant_prime_mover" plant part record and a "plant_unit" plant part record
which were both cobbled together from the same two generators. We want to be
able to reduce the plant-part list to only unique collections of generators,
so we label the first unique granularity as a true granularity and label the
subsequent records as false granularities with the ``true_gran`` column. In
order to choose which plant-part to keep in these instances, we assigned a
:py:const:`PLANT_PARTS_ORDERED` and label whichever plant-part comes first as
the unique granularity.

**Recipe Book for the plant-part list**

:py:const:`PLANT_PARTS` is the main recipe book for how each of the plant-parts
need to be compiled. These plant-parts represent ways to group generators based
on widely reported values in EIA. All of these are logical ways to group
collections of generators - in most cases - but some groupings of generators
are more prevelant or relevant than others for certain types of plants.

The canonical example here is the ``plant_unit``. A unit is a collection of
generators that operate together - most notably the combined-cycle natural gas
plants. Combined-cycle units generally consist of a number of gas turbines
which feed excess steam to a number of steam turbines.

>>> df_gens = pd.DataFrame({
...     'plant_id_eia': [1, 1, 1],
...     'generator_id': ['a', 'b', 'c'],
...     'unit_id_pudl': [1, 1, 1],
...     'prime_mover_code': ['CT', 'CT', 'CA'],
...     'capacity_mw': [50, 50, 100],
... })
>>> df_gens
    plant_id_eia    generator_id    unit_id_pudl    prime_mover_code    capacity_mw
0              1               a               1                  CT             50
1              1               b               1                  CT             50
2              1               c               1                  CA            100

A good example of a plant-part that isn't really logical also comes from a
combined-cycle unit. Grouping this example plant by the ``prime_mover_code``
would generate two records that would basically never show up in FERC1.
This stems from the inseparability of the generators.

>>> df_plant_prime_mover = pd.DataFrame({
...     'plant_id_eia': [1, 1],
...     'plant_part': ['plant_prime_mover', 'plant_prime_mover'],
...     'prime_mover_code': ['CT', 'CA'],
...     'capacity_mw': [100, 100],
... })
>>> df_plant_prime_mover
    plant_id_eia         plant_part    prime_mover_code    capacity_mw
0              1  plant_prime_mover                  CT            100
1              1  plant_prime_mover                  CA            100

In this case the unit is more relevant:

>>> df_plant_unit = pd.DataFrame({
...     'plant_id_eia': [1],
...     'plant_part': ['plant_unit'],
...     'unit_id_pudl': [1],
...     'capacity_mw': [200],
... })
>>> df_plant_unit
    plant_id_eia    plant_part    unit_id_pudl    capacity_mw
0              1    plant_unit               1            200

But if this same plant had both this combined-cycle unit and two more
generators that were self contained "GT" or gas combustion turbine, a frequent
way to group these generators is differnt for the combined-cycle unit and the
gas-turbine.

>>> df_gens = pd.DataFrame({
...     'plant_id_eia': [1, 1, 1, 1, 1],
...     'generator_id': ['a', 'b', 'c', 'd', 'e'],
...     'unit_id_pudl': [1, 1, 1, 2, 3],
...     'prime_mover_code': ['CT', 'CT', 'CA', 'GT', 'GT'],
...     'capacity_mw': [50, 50, 100, 75, 75],
... })
>>> df_gens
    plant_id_eia    generator_id    unit_id_pudl    prime_mover_code    capacity_mw
0              1               a               1                  CT             50
1              1               b               1                  CT             50
2              1               c               1                  CA            100
3              1               d               2                  GT             75
4              1               e               3                  GT             75

>>> df_plant_part = pd.DataFrame({
...     'plant_id_eia': [1, 1],
...     'plant_part': ['plant_unit', 'plant_prime_mover'],
...     'unit_id_pudl': [1, pd.NA],
...     'prime_mover_code': [pd.NA, 'GT',],
...     'capacity_mw': [200, 150],
... })
>>> df_plant_part
    plant_id_eia           plant_part    unit_id_pudl    prime_mover_code    capacity_mw
0              1           plant_unit               1                <NA>            200
1              1    plant_prime_mover            <NA>                  GT            150

In this case last, the ``plant_unit`` record would have a null
``plant_prime_mover`` because the unit contains more than one
``prime_mover_code``. Same goes for the ``unit_id_pudl`` of the
``plant_prime_mover``. This is handled in the :class:``AddConsistentAttributes``.

**Overview of flow for generating the master unit list:**

The three main classes which enable the generation of the plant-part table are:

* :class:`MakeMegaGenTbl`: All of the plant parts are compiled from generators.
  So this class generates a big dataframe of generators with any ID and data
  columns we'll need. This is also where we add records regarding utility
  ownership slices. The table includes two records for every generator-owner:
  one for the "total" generator (assuming the owner owns 100% of the generator)
  and one for the report ownership fraction of that generator with all of the
  data columns scaled to the ownership fraction.
* :class:`LabelTrueGranularities`: This class creates labels for all generators
  which note wether the plant-part records that will be compiled from each
  generator will be a "true granulary", as described above.
* :class:`MakePlantParts`: This class uses the generator dataframe, the
  granularity dataframe from the above two classes as well as the information
  stored in :py:const:`PLANT_PARTS` to know how to aggregate each of the plant
  parts. Then we have plant part dataframes with the columns which identify the
  plant part and all of the data columns aggregated to the level of the plant
  part. With that compiled plant part dataframe we also add in qualifier
  columns with :class:`AddConsistentAttributes`. A qualifer column is a column which
  contain data that is not endemic to the plant part record (it is not one of
  the identifying columns or aggregated data columns) but the data is still
  useful data that is attributable to each of the plant part records. For more
  detail on what a qualifier column is, see :meth:`AddConsistentAttributes.execute`.

**Generating the plant-parts list**

There are two ways to generate the plant-parts table: one directly using the
:class:`pudl.output.pudltabl.PudlTabl` object and the other using the classes
from this module. Either option needs a :class:`pudl.output.pudltabl.PudlTabl`
object.

Create the :class:`pudl.output.pudltabl.PudlTabl` object:

.. code-block:: python

    import pudl
    pudl_engine = sa.create_engine(pudl.workspace.setup.get_defaults()['pudl_db'])
    pudl_out = pudl.output.pudltabl.PudlTabl(pudl_engine,freq='AS')

Then make the table via pudl_out:

.. code-block:: python

    plant_parts_eia = pudl_out.plant_parts_eia()


OR make the table via objects in this module:

.. code-block:: python

    gens_mega = MakeMegaGenTbl().execute(mcoe, own_eia860)
    true_grans = LabelTrueGranularities().execute(gens_mega)
    parts_compiler = MakePlantParts(pudl_out)
    plant_parts_eia = parts_compiler.execute(gens_mega=gens_mega, true_grans=true_grans)


"""

import logging
import warnings
from copy import deepcopy
from typing import Dict, List, Literal

import numpy as np
import pandas as pd

import pudl

logger = logging.getLogger(__name__)

# HALP: I need both of these setting set in order for the dfs in the docstrings
# to pass the doctests. Without them the formatting get all jumbled.
# but obviously this is the wrong place to do this.
# I tried adding these into conftest.py in pandas_terminal_width().
# I tried adding this into __init__.py.
# I tried adding this into the module docstring.
pd.options.display.width = 1000
pd.options.display.max_columns = 1000

PLANT_PARTS: Dict[str, Dict[str, List]] = {
    "plant": {
        "id_cols": ["plant_id_eia"],
    },
    "plant_gen": {
        "id_cols": ["plant_id_eia", "generator_id"],
    },
    "plant_unit": {
        "id_cols": ["plant_id_eia", "unit_id_pudl"],
    },
    "plant_technology": {
        "id_cols": ["plant_id_eia", "technology_description"],
    },
    "plant_prime_fuel": {  # 'plant_primary_fuel': {
        "id_cols": ["plant_id_eia", "energy_source_code_1"],
    },
    "plant_prime_mover": {
        "id_cols": ["plant_id_eia", "prime_mover_code"],
    },
    "plant_ferc_acct": {
        "id_cols": ["plant_id_eia", "ferc_acct_name"],
    },
}
"""
dict: this dictionary contains a key for each of the 'plant parts' that should
end up in the mater unit list. The top-level value for each key is another
dictionary, which contains keys:

* id_cols (the primary key type id columns for this plant part). The
  plant_id_eia column must come first.

"""

PLANT_PARTS_ORDERED: List[str] = [
    "plant",
    "plant_unit",
    "plant_prime_mover",
    "plant_technology",
    "plant_prime_fuel",
    "plant_ferc_acct",
    "plant_gen",
]


IDX_TO_ADD: List[str] = ["report_date", "operational_status_pudl"]
"""
list: list of additional columns to add to the id_cols in :py:const:`PLANT_PARTS`.
The id_cols are the base columns that we need to aggregate on, but we also need
to add the report date to keep the records time sensitive and the
operational_status_pudl to separate the operating plant-parts from the
non-operating plant-parts.
"""

IDX_OWN_TO_ADD: List[str] = ["utility_id_eia", "ownership"]
"""
list: list of additional columns beyond the :py:const:`IDX_TO_ADD` to add to the
id_cols in :py:const:`PLANT_PARTS` when we are dealing with plant-part records
that have been broken out into "owned" and "total" records for each of their
owners.
"""

SUM_COLS: List[str] = [
    "total_fuel_cost",
    "net_generation_mwh",
    "capacity_mw",
    "capacity_eoy_mw",
    "total_mmbtu",
]
"""
list: list of columns to sum when aggregating a table.
"""

WTAVG_DICT = {
    "fuel_cost_per_mwh": "capacity_mw",
    "heat_rate_mmbtu_mwh": "capacity_mw",
    "fuel_cost_per_mmbtu": "capacity_mw",
}
"""
dict: a dictionary of columns (keys) to perform weighted averages on and
the weight column (values)"""


CONSISTENT_ATTRIBUTE_COLS = [
    "fuel_type_code_pudl",
    "planned_retirement_date",
    "retirement_date",
    "generator_id",
    "unit_id_pudl",
    "technology_description",
    "energy_source_code_1",
    "prime_mover_code",
    "ferc_acct_name",
]
"""
list: a list of column names to add as attributes when they are consistent into
the aggregated plant-part records.
"""

PRIORITY_ATTRIBUTES_DICT = {
    "operational_status": ["existing", "proposed", "retired"],
}


FIRST_COLS = [
    "plant_id_eia",
    "report_date",
    "plant_part",
    "generator_id",
    "unit_id_pudl",
    "prime_mover_code",
    "energy_source_code_1",
    "technology_description",
    "ferc_acct_name",
    "utility_id_eia",
    "true_gran",
    "appro_part_label",
]


class MakeMegaGenTbl(object):
    """
    Compiler for a MEGA generator table with ownership integrated.

    Examples
    --------
    **Input Tables**

    Here is an example of one plant with three generators. We will use
    ``capacity_mw`` as the data column.

    >>> mcoe = pd.DataFrame({
    ...     'plant_id_eia': [1, 1, 1],
    ...     'report_date': ['2020-01-01', '2020-01-01','2020-01-01'],
    ...     'generator_id': ['a', 'b', 'c'],
    ...     'utility_id_eia': [111, 111, 111],
    ...     'unit_id_pudl': [1, 1, 1],
    ...     'prime_mover_code': ['CT', 'CT', 'CA'],
    ...     'technology_description': [
    ...         'Natural Gas Fired Combined Cycle', 'Natural Gas Fired Combined Cycle', 'Natural Gas Fired Combined Cycle'
    ...     ],
    ...     'operational_status': ['existing', 'existing','existing'],
    ...     'retirement_date': [pd.NA, pd.NA, pd.NA],
    ...     'capacity_mw': [50, 50, 100],
    ... }).astype({
    ...     'retirement_date': "datetime64[ns]",
    ...     'report_date': "datetime64[ns]",
    ... })
    >>> mcoe
        plant_id_eia    report_date 	generator_id   utility_id_eia 	unit_id_pudl 	prime_mover_code              technology_description   operational_status  retirement_date 	capacity_mw
    0 	           1     2020-01-01 	           a              111 	           1 	              CT    Natural Gas Fired Combined Cycle             existing               NaT 	         50
    1 	           1     2020-01-01 	           b              111 	           1 	              CT    Natural Gas Fired Combined Cycle             existing               NaT 	         50
    2 	           1     2020-01-01 	           c              111 	           1 	              CA    Natural Gas Fired Combined Cycle             existing               NaT 	        100

    The ownership table from EIA 860 includes one record for every owner of
    each generator. In this example generator ``c`` has two owners.

    >>> df_own_eia860 = pd.DataFrame({
    ...     'plant_id_eia': [1, 1, 1, 1],
    ...     'report_date': ['2020-01-01', '2020-01-01','2020-01-01', '2020-01-01'],
    ...     'generator_id': ['a', 'b', 'c', 'c'],
    ...     'utility_id_eia': [111, 111, 111, 111],
    ...     'owner_utility_id_eia': [111, 111, 111, 888],
    ...     'fraction_owned': [1, 1, .75, .25]
    ... }).astype({'report_date': "datetime64[ns]"})
    >>> df_own_eia860
        plant_id_eia 	report_date   generator_id 	utility_id_eia 	owner_utility_id_eia  fraction_owned
    0 	           1 	 2020-01-01 	         a 	           111 	                 111 	        1.00
    1 	           1 	 2020-01-01 	         b 	           111 	                 111 	        1.00
    2 	           1 	 2020-01-01 	         c 	           111 	                 111 	        0.75
    3 	           1 	 2020-01-01 	         c 	           111 	                 888 	        0.25

    **Output Mega Generators Table**

    ``MakeMegaGenTbl().execute(mcoe, df_own_eia860, slice_cols=['capacity_mw'])``
    produces the output table ``gens_mega`` which includes two main sections:
    the generators with a "total" ownership stake for each of their owners and
    the generators with an "owned" ownership stake for each of their
    owners. For the generators that are owned 100% by one utility, the
    records are identical except the ``ownership`` column. For the
    generators that have more than one owner, there are two "total" records
    with 100% of the capacity of that generator - one for each owner - and
    two "owned" records with the capacity scaled to the ownership stake
    of each of the owner utilites - represented by ``fraction_owned``.

    """

    def __init__(self):
        """
        Initialize object which creates a MEGA generator table.

        The coordinating function here is :meth:`execute`.

        """
        self.id_cols_list = make_id_cols_list()

    def execute(
        self,
        mcoe: pd.DataFrame,
        own_eia860: pd.DataFrame,
        slice_cols: List[str] = SUM_COLS,
        validate_own_merge: str = "1:m",
    ) -> pd.DataFrame:
        """
        Make the mega generators table with ownership integrated.

        Args:
            mcoe: generator-based mcoe table from :meth:`pudl.output.PudlTabl.mcoe()`
            own_eia860: ownership table from :meth:`pudl.output.PudlTabl.own_eia860()`
            slice_cols: list of columns to slice by ownership fraction in
                :meth:`MakeMegaGenTbl.slice_by_ownership`. Default is :py:const:`SUM_COLS`
            validate_own_merge: how the merge between ``mcoe`` and ``owm_eia860``
                is to be validated via ``pd.merge``. If there should be one
                record for each plant/generator/date in ``mcoe`` then the default
                `1:m` should be used.

        Returns:
            a table of all of the generators with identifying columns and data
            columns, sliced by ownership which makes "total" and "owned"
            records for each generator owner. The "owned" records have the
            generator's data scaled to the ownership percentage (e.g. if a 200
            MW generator has a 75% stake owner and a 25% stake owner, this will
            result in two "owned" records with 150 MW and 50 MW). The "total"
            records correspond to the full plant for every owner (e.g. using
            the same 2-owner 200 MW generator as above, each owner will have a
            records with 200 MW).

        """
        logger.info("Generating the mega generator table with ownership.")

        gens_mega = (
            self.get_gens_mega_table(mcoe)
            .pipe(self.label_operating_gens)
            .pipe(self.slice_by_ownership, own_eia860, slice_cols, validate_own_merge)
        )
        return gens_mega

    def get_gens_mega_table(self, mcoe):
        """
        Compile the main generators table that will be used as base of PPL.
{'fraction_non_alphanumeric': 0.09177605143316367, 'fraction_numerical': 0.02397535494240557, 'mean_word_length': 2.5082229113805092, 'pattern_counts': {'":': 36, '<': 7, '<?xml version=': 0, '>': 93, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 66, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '232456', 'n_tokens_mistral': 11996, 'n_tokens_neox': 10915, 'n_words': 4084}