8def gcor(
  9    x,
 10    y=None,
 11    *,
 12    drop_na='none',
 13    k=None,
 14    max_levels=100,
 15) -> Union[float, pd.DataFrame]:
 16    """
 17    Generalized correlation measure.
 18
 19    Parameters
 20    ----------
 21    x : 1d array-like or `pandas.DataFrame`
 22        Input data. If 1d array-like, treated as n observations of a random variable.
 23        If DataFrame, treated as n observations of p random variables (columns).
 24        Must not be None.
 25    
 26    y : 1d array-like or None, default None
 27        Second input data. Required if `x` is 1d array-like.
 28        It is treated as n observations of another random variable.
 29        Must be None if `x` is a DataFrame.
 30    
 31    drop_na : {'none', 'pairwise', 'casewise'}, default 'none'
 32        How to handle missing values:
 33        - 'none': Keep missing values and treat them as observations of a separate category
 34        - 'pairwise': For each variable pair (X, Y), drop the i-th observation
 35          if either X or Y is missing at position i
 36        - 'casewise': Drop the i-th observation for all variables,
 37          if any variable is missing at position i
 38    
 39    k : int or None, default None
 40        Number of quantile divisions. If None, determined automatically.
 41    
 42    max_levels : int, default 100
 43        The maximum number of levels allowed when converting non-numeric variables to categories.
 44
 45    Returns
 46    -------
 47    float or pandas.DataFrame
 48        If `x` is 1d array-like, returns a float (or NaN).
 49        If `x` is DataFrame, returns a generalized correlation matrix as a DataFrame.
 50        The diagonal elements are always set to 1.0.
 51    
 52    Raises
 53    ------
 54    ValueError
 55        If `drop_na` is invalid or if the lengths of `x` and `y` do not match.
 56    TypeError
 57        If the combination of `x` and `y` is invalid.
 58    
 59    Examples
 60    --------
 61    >>> from gcor import gcor
 62    >>> x = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
 63    >>> y = [1, 2, 3, 4, 5, 3, 4, 5, 6, 7]
 64    >>> gcor(x, y)
 65    0.5345224838248488
 66
 67    >>> import pandas as pd
 68    >>> df = pd.DataFrame({
 69    ...     'x': x,
 70    ...     'y': y,
 71    ...     'z': ['a', 'a', 'b', 'b', 'c', 'c', 'd', 'd', 'e', 'e'],
 72    ... })
 73    >>> gcor(df)
 74              x         y         z
 75    x  1.000000  0.534522  0.838289
 76    y  0.534522  1.000000  0.763233
 77    z  0.838289  0.763233  1.000000
 78    """
 79    # --- Validation ---
 80    valid_drop_na = {'none', 'pairwise', 'casewise'}
 81    if drop_na not in valid_drop_na:
 82        raise ValueError(
 83            "drop_na must be one of {'none', 'pairwise', 'casewise'} "
 84            f"(got {drop_na!r})."
 85        )
 86    
 87    if x is None:
 88        raise TypeError("x must not be None.")
 89
 90    # --- Store inputs in a temporary DataFrame (df_tmp) ---
 91    return_matrix: bool = isinstance(x, pd.DataFrame)
 92
 93    if return_matrix:
 94        if y is not None:
 95            raise TypeError('y must be None when x is a pandas.DataFrame.')
 96        df_tmp = x.copy(deep=False)
 97    else:
 98        if y is None:
 99            raise TypeError('y must not be None when x is not a pandas.DataFrame.')
100
101        xs = None
102        ys = None
103
104        # If both are Series, align by position and warn when indices differ.
105        if isinstance(x, pd.Series) and isinstance(y, pd.Series):
106            if not x.index.equals(y.index):
107                warnings.warn(
108                    'x and y are pandas.Series but their indices do not match. '
109                    'Resetting indices (drop=True) and matching by position.',
110                    UserWarning,
111                    stacklevel=2,
112                )
113                xs = x.reset_index(drop=True)
114                ys = y.reset_index(drop=True)
115
116        # If xs/ys are still None, convert inputs to Series (or keep as-is if already Series).
117        if xs is None:
118            xs = x if isinstance(x, pd.Series) else pd.Series(x)
119        if ys is None:
120            ys = y if isinstance(y, pd.Series) else pd.Series(y)
121
122        # Ensure lengths match.
123        if len(xs) != len(ys):
124            raise ValueError(
125                f'x and y must have the same length (got len(x)={len(xs)} and len(y)={len(ys)}).'
126            )
127
128        df_tmp = pd.DataFrame({'x': xs, 'y': ys})
129
130    # --- Missing-value handling (casewise) ---
131    if drop_na == 'casewise':
132        df_tmp = df_tmp.dropna(how='any')
133
134    # --- Discretize each column ---
135    df = df_tmp.apply(lambda s: discretize(s, k=k, max_levels=max_levels), axis=0)
136
137    if len(df) > 0:
138        non_categorical_cols = [
139            col for col in df.columns
140            if not isinstance(df[col].dtype, pd.CategoricalDtype)
141        ]
142        if non_categorical_cols:
143            warnings.warn(
144                "The following columns were not converted to categorical and may "
145                "result in NaN values: "
146                + ", ".join(map(str, non_categorical_cols))
147                + ". Try setting the 'max_levels' argument if they have too many levels.",
148                UserWarning,
149                stacklevel=2,
150            )
151
152    # Precompute missing-value mask only when needed (pairwise).
153    if drop_na == 'pairwise':
154        na = df.isna()
155
156    # --- Compute generalized correlation matrix ---
157    ncol = df.shape[1]
158    mat = np.full((ncol, ncol), np.nan, dtype=float)
159
160    for i in range(ncol):
161        for j in range(i, ncol):
162            if i == j:
163                mat[i, j] = 1.0
164                continue
165
166            # Missing value handling (pairwise)
167            if drop_na == 'pairwise':
168                mask = ~(na.iloc[:, i] | na.iloc[:, j])
169                val = gcor_cat(df.iloc[mask, i], df.iloc[mask, j])
170            else:
171                val = gcor_cat(df.iloc[:, i], df.iloc[:, j])
172
173            mat[i, j] = val
174            mat[j, i] = val
175
176    if return_matrix:
177        return pd.DataFrame(mat, index=df.columns, columns=df.columns)
178    else:
179        # Return as Python float (cast from NumPy scalar)
180        return float(mat[0, 1])