o
    Gj\"                     @   s\   d Z ddlmZmZmZ ddlmZ ddlmZ ddl	m
Z
 ddlmZ G dd deZd	S )
z
Regression Controller
Handles HTTP layer for ML model training, scoring, and prediction.
Follows BaseController(ApiResponseMixin) pattern from prophet_controller.py.
    )DictAnyOptional)BaseController)RegressionService)OMSRegressionDataService)
get_loggerc                       s   e Zd ZdZ fddZdeeef defddZ	deeef defdd	Z
deeef defd
dZdeeef defddZdedefddZdefddZdedefddZ  ZS )RegressionControllerz@Controller for XGBoost regression and classification operations.c                    s(   t    t | _t | _td| _d S )Nregression_controller)super__init__r   servicer   data_servicer   logger)self	__class__ C/var/www/html/oms-prod-mcp/app/controllers/regression_controller.pyr      s   
zRegressionController.__init__datareturnc           
   
   C   s  z^t |dr
| }|d}|s| jddW S d| }| jt|}|js.t|dk r>| j	d| dt| d	dW S | jj
}|| }|d
 }| jj|||d|dd}| |dW S  ty }	 z| jjdt|	 dd | jdt|	 dW  Y d}	~	S d}	~	ww )z/Train collection risk classifier for a company.
model_dump
company_idcompany_id is requiredmessagecollection_risk_co   'Insufficient training data for company u&    (need ≥5 clients with history, got )is_high_risk
classifierconfigmodel_idXy
model_typer"   z*Collection risk model trained successfullyztrain_collection_risk error: Texc_infoTraining failed: N)hasattrr   getbad_requestr   !get_collection_risk_training_dataintemptylenunprocessable_entityCOLLECTION_FEATURESr   traincreated	Exceptionr   errorstrserver_error)
r   r   r   r$   dffeaturesr%   r&   resulter   r   r   train_collection_risk   s<   


"z*RegressionController.train_collection_riskc           	   
      s   zet |dr
| }|d}|s| jddW S d| }| jt|  js.t dk r>| j	d| dt  d	dW S | jj
}  fd
d|D  } d }| jj|||d|dd}| |dW S  ty } z| jjdt| dd | jdt| dW  Y d}~S d}~ww )z6Train sales target attainment regressor for a company.r   r   r   r   sales_target_cor   r   u'    (need ≥5 district×season rows, got r   c                       g | ]	}| j v r|qS r   columns.0fr:   r   r   
<listcomp>T       z;RegressionController.train_sales_target.<locals>.<listcomp>attainment_pct	regressorr"   r#   z'Sales target model trained successfullyztrain_sales_target error: Tr(   r*   N)r+   r   r,   r-   r   get_sales_target_training_datar/   r0   r1   r2   SALES_TARGET_FEATURESr   r4   r5   r6   r   r7   r8   r9   )	r   r   r   r$   r;   r%   r&   r<   r=   r   rF   r   train_sales_target@   s<   


"z'RegressionController.train_sales_targetc              
   C   sn  zst |dr
| }|d}|s| jddW S |dpd}d| }| jt||}|jr8| jg ddW S | j	
||}g }t| D ]\}\}	}
tt|| d	}|t|
d
 |t|d qG| |dt| dW S  ty } z| jt|dW  Y d}~S d}~w ty } z| jjdt| dd | jdt| dW  Y d}~S d}~ww )z0Score clients for collection payment delay risk.r   r   r   r   
client_idsNr   zNo scoreable clients found   	client_id)rP   delay_probability
risk_levelzScored z clientszscore_collection_risk error: Tr(   zScoring failed: )r+   r   r,   r-   r   $get_collection_risk_scoring_featuresr/   r0   okr   predict_proba	enumerateiterrowsroundfloatappendr   collection_risk_levelr1   
ValueErrorr2   r8   r6   r   r7   r9   )r   r   r   rN   r$   features_dfprobabilitiesresultsi_rowprobr=   r   r   r   score_collection_riskh   s>   




"z*RegressionController.score_collection_riskc              
      s  zt |dr
| }|d}|d}|r|s| jddW S d| }| jjt|t||d|dd	  jrC| jg d
dW S | jj	}  fdd|D  }| j
||}g }t  D ])\}	\}
}tt||	 d}|t|d t|d t|d |t|d qa| |dt| dW S  ty } z| jt|dW  Y d}~S d}~w ty } z| jjdt| dd | jdt| dW  Y d}~S d}~ww )u=   Predict attainment % for district×product_head combinations.r   r   account_year_idz+company_id and account_year_id are requiredr   r?   product_head_idsdistrict_ids)rf   rg   z1No prediction targets found for this account yearc                    r@   r   rA   rC   r]   r   r   rG      rH   z=RegressionController.predict_sales_target.<locals>.<listcomp>   product_head_iddistrict_id)rj   rk   re   predicted_attainment_pctrR   z
Predicted z combinationsNzpredict_sales_target error: Tr(   zPrediction failed: )r+   r   r,   r-   r   $get_sales_target_prediction_featuresr/   r0   rT   rL   r   predictrV   rW   rX   rY   rZ   r   sales_target_risk_levelr1   r\   r2   r8   r6   r   r7   r9   )r   r   r   re   r$   feature_colsr%   predictionsr_   r`   ra   rb   pctr=   r   rh   r   predict_sales_target   sL   







"z)RegressionController.predict_sales_targetr$   c              
   C   sl   z| j |}|d u r| jd| ddW S | |W S  ty5 } z| jt|dW  Y d }~S d }~ww )NModel '' not foundr   )r   get_model_info	not_foundrT   r6   r9   r8   )r   r$   infor=   r   r   r   rv      s   z#RegressionController.get_model_infoc              
   C   sX   z| j  }| |t| dW S  ty+ } z| jt|dW  Y d }~S d }~ww )Nz regression models foundr   )r   list_modelsrT   r1   r6   r9   r8   )r   modelsr=   r   r   r   ry      s   
z RegressionController.list_modelsc              
   C   sn   z| j |}|s| jd| ddW S | d|idW S  ty6 } z| jt|dW  Y d }~S d }~ww )Nrt   ru   r   r$   zModel deleted)r   delete_modelrw   rT   r6   r9   r8   )r   r$   deletedr=   r   r   r   r{      s   z!RegressionController.delete_model)__name__
__module____qualname____doc__r   r   r8   r   tupler>   rM   rd   rs   rv   ry   r{   __classcell__r   r   r   r   r	      s    
$(&0	r	   N)r   typingr   r   r   app.controllers.base_controllerr   app.services.regression_servicer   (app.services.oms_regression_data_servicer   app.utils.helpersr   r	   r   r   r   r   <module>   s    