import React, { useState } from 'react';

const ProbabilityTreeWidget: React.FC = () => {
  const [replacement, setReplacement] = useState(false);
  const [initR, setInitR] = useState(3);
  const [initB, setInitB] = useState(4);
  const [hoverPath, setHoverPath] = useState<string | null>(null); // 'RR', 'RB', 'BR', 'BB'

  const total = initR + initB;

  // Level 1 Probs
  const pR = initR / total;
  const pB = initB / total;

  // Level 2 Probs (Given R)
  const r_R = replacement ? initR : Math.max(0, initR - 1);
  const r_Total = replacement ? total : total - 1;
  const pR_R = r_Total > 0 ? r_R / r_Total : 0;
  const pB_R = r_Total > 0 ? 1 - pR_R : 0;

  // Level 2 Probs (Given B)
  const b_B = replacement ? initB : Math.max(0, initB - 1);
  const b_Total = replacement ? total : total - 1;
  const pB_B = b_Total > 0 ? b_B / b_Total : 0;
  const pR_B = b_Total > 0 ? 1 - pB_B : 0;

  // Final Probs
  const pRR = pR * pR_R;
  const pRB = pR * pB_R;
  const pBR = pB * pR_B;
  const pBB = pB * pB_B;

  const fraction = (num: number, den: number) => {
      if (den === 0) return "0";
      return (
        <span className="font-mono bg-white px-1 rounded shadow-sm border border-slate-200 text-xs inline-block mx-1">
            {num}/{den}
        </span>
      );
  };

  const getPathColor = (path: string, segment: 'top' | 'bottom' | 'top-top' | 'top-bottom' | 'bottom-top' | 'bottom-bottom') => {
      const defaultColor = "#cbd5e1"; // Slate 300
      
      if (!hoverPath) {
          // Default coloring based on branch type
          if (segment.includes('top')) return "#f43f5e"; // Red branches
          if (segment.includes('bottom')) return "#3b82f6"; // Blue branches
          return defaultColor;
      }

      // Highlighting logic based on hoverPath
      if (segment === 'top') return (hoverPath === 'RR' || hoverPath === 'RB') ? "#f43f5e" : "#f1f5f9";
      if (segment === 'bottom') return (hoverPath === 'BR' || hoverPath === 'BB') ? "#3b82f6" : "#f1f5f9";
      
      if (segment === 'top-top') return hoverPath === 'RR' ? "#f43f5e" : "#f1f5f9";
      if (segment === 'top-bottom') return hoverPath === 'RB' ? "#3b82f6" : "#f1f5f9";
      
      if (segment === 'bottom-top') return hoverPath === 'BR' ? "#f43f5e" : "#f1f5f9";
      if (segment === 'bottom-bottom') return hoverPath === 'BB' ? "#3b82f6" : "#f1f5f9";

      return defaultColor;
  };
  
  const getStrokeWidth = (segment: string) => {
      if (!hoverPath) return 2;
      
      if (segment === 'top') return (hoverPath === 'RR' || hoverPath === 'RB') ? 4 : 1;
      if (segment === 'bottom') return (hoverPath === 'BR' || hoverPath === 'BB') ? 4 : 1;
      
      if (segment === 'top-top') return hoverPath === 'RR' ? 4 : 1;
      if (segment === 'top-bottom') return hoverPath === 'RB' ? 4 : 1;
      
      if (segment === 'bottom-top') return hoverPath === 'BR' ? 4 : 1;
      if (segment === 'bottom-bottom') return hoverPath === 'BB' ? 4 : 1;
      
      return 2;
  }

  return (
    <div className="bg-white p-6 rounded-xl shadow-lg border border-slate-200">
       
       {/* Controls */}
       <div className="flex flex-wrap justify-between items-center mb-6 gap-4">
           <div className="flex gap-4">
               <div className="flex flex-col">
                   <label className="text-xs font-bold text-rose-600 uppercase mb-1">Red Items</label>
                   <div className="flex items-center gap-2">
                       <button onClick={() => setInitR(Math.max(1, initR-1))} className="w-6 h-6 bg-rose-100 text-rose-700 rounded hover:bg-rose-200">-</button>
                       <span className="font-bold w-4 text-center">{initR}</span>
                       <button onClick={() => setInitR(Math.min(10, initR+1))} className="w-6 h-6 bg-rose-100 text-rose-700 rounded hover:bg-rose-200">+</button>
                   </div>
               </div>
               <div className="flex flex-col">
                   <label className="text-xs font-bold text-blue-600 uppercase mb-1">Blue Items</label>
                   <div className="flex items-center gap-2">
                       <button onClick={() => setInitB(Math.max(1, initB-1))} className="w-6 h-6 bg-blue-100 text-blue-700 rounded hover:bg-blue-200">-</button>
                       <span className="font-bold w-4 text-center">{initB}</span>
                       <button onClick={() => setInitB(Math.min(10, initB+1))} className="w-6 h-6 bg-blue-100 text-blue-700 rounded hover:bg-blue-200">+</button>
                   </div>
               </div>
           </div>

           <div className="flex bg-slate-100 p-1 rounded-lg">
               <button 
                 onClick={() => setReplacement(true)}
                 className={`px-3 py-1 text-xs font-bold rounded transition-all ${replacement ? 'bg-white shadow text-indigo-600' : 'text-slate-500'}`}
               >
                   With Replacement
               </button>
               <button 
                 onClick={() => setReplacement(false)}
                 className={`px-3 py-1 text-xs font-bold rounded transition-all ${!replacement ? 'bg-white shadow text-indigo-600' : 'text-slate-500'}`}
               >
                   Without Replacement
               </button>
           </div>
       </div>

       <div className="relative h-64 w-full max-w-lg mx-auto select-none">
           <svg width="100%" height="100%" className="overflow-visible">
               {/* Root */}
               <circle cx="20" cy="128" r="6" fill="#64748b" />
               
               {/* Level 1 Branches */}
               <path d="M 20 128 C 50 128, 50 64, 150 64" fill="none" stroke={getPathColor('R', 'top')} strokeWidth={getStrokeWidth('top')} className="transition-all duration-300" />
               <path d="M 20 128 C 50 128, 50 192, 150 192" fill="none" stroke={getPathColor('B', 'bottom')} strokeWidth={getStrokeWidth('bottom')} className="transition-all duration-300" />
               
               {/* Level 1 Labels */}
               <foreignObject x="60" y="70" width="60" height="30">
                   <div className={`text-center font-bold text-xs ${hoverPath && (hoverPath[0]!=='R') ? 'text-slate-300' : 'text-rose-600'}`}>{initR}/{total}</div>
               </foreignObject>
               <foreignObject x="60" y="150" width="60" height="30">
                   <div className={`text-center font-bold text-xs ${hoverPath && (hoverPath[0]!=='B') ? 'text-slate-300' : 'text-blue-600'}`}>{initB}/{total}</div>
               </foreignObject>

               {/* Level 1 Nodes */}
               <circle cx="150" cy="64" r="18" fill="#f43f5e" className={`transition-all ${hoverPath && hoverPath[0] !== 'R' ? 'opacity-20' : 'opacity-100 shadow-md'}`} />
               <text x="150" y="68" textAnchor="middle" fill="white" className={`text-xs font-bold pointer-events-none ${hoverPath && hoverPath[0] !== 'R' ? 'opacity-20' : ''}`}>R</text>

               <circle cx="150" cy="192" r="18" fill="#3b82f6" className={`transition-all ${hoverPath && hoverPath[0] !== 'B' ? 'opacity-20' : 'opacity-100 shadow-md'}`} />
               <text x="150" y="196" textAnchor="middle" fill="white" className={`text-xs font-bold pointer-events-none ${hoverPath && hoverPath[0] !== 'B' ? 'opacity-20' : ''}`}>B</text>

               {/* Level 2 Branches (Top) */}
               <path d="M 168 64 L 280 32" fill="none" stroke={getPathColor('RR', 'top-top')} strokeWidth={getStrokeWidth('top-top')} strokeDasharray="4,2" className="transition-all duration-300" />
               <path d="M 168 64 L 280 96" fill="none" stroke={getPathColor('RB', 'top-bottom')} strokeWidth={getStrokeWidth('top-bottom')} strokeDasharray="4,2" className="transition-all duration-300" />
               
               {/* Level 2 Top Labels */}
               <foreignObject x="190" y="25" width="60" height="30">
                   <div className={`text-center font-bold text-xs ${hoverPath === 'RR' ? 'text-rose-600 scale-110' : 'text-slate-400'}`}>{r_R}/{r_Total}</div>
               </foreignObject>
               <foreignObject x="190" y="80" width="60" height="30">
                   <div className={`text-center font-bold text-xs ${hoverPath === 'RB' ? 'text-blue-600 scale-110' : 'text-slate-400'}`}>{initB}/{r_Total}</div>
               </foreignObject>

               {/* Level 2 Branches (Bottom) */}
               <path d="M 168 192 L 280 160" fill="none" stroke={getPathColor('BR', 'bottom-top')} strokeWidth={getStrokeWidth('bottom-top')} strokeDasharray="4,2" className="transition-all duration-300" />
               <path d="M 168 192 L 280 224" fill="none" stroke={getPathColor('BB', 'bottom-bottom')} strokeWidth={getStrokeWidth('bottom-bottom')} strokeDasharray="4,2" className="transition-all duration-300" />

               {/* Level 2 Bottom Labels */}
               <foreignObject x="190" y="150" width="60" height="30">
                   <div className={`text-center font-bold text-xs ${hoverPath === 'BR' ? 'text-rose-600 scale-110' : 'text-slate-400'}`}>{initR}/{b_Total}</div>
               </foreignObject>
               <foreignObject x="190" y="210" width="60" height="30">
                   <div className={`text-center font-bold text-xs ${hoverPath === 'BB' ? 'text-blue-600 scale-110' : 'text-slate-400'}`}>{b_B}/{b_Total}</div>
               </foreignObject>

               {/* Outcomes (Interactive Targets) */}
               <g 
                 className="cursor-pointer" 
                 onMouseEnter={() => setHoverPath('RR')} 
                 onMouseLeave={() => setHoverPath(null)}
               >
                   <text x="300" y="36" className={`text-xs font-bold transition-all ${hoverPath === 'RR' ? 'fill-rose-600 text-base' : 'fill-slate-500'}`}>RR: {(pRR * 100).toFixed(1)}%</text>
                   <rect x="290" y="20" width="80" height="20" fill="transparent" />
               </g>

               <g 
                 className="cursor-pointer" 
                 onMouseEnter={() => setHoverPath('RB')} 
                 onMouseLeave={() => setHoverPath(null)}
               >
                   <text x="300" y="100" className={`text-xs font-bold transition-all ${hoverPath === 'RB' ? 'fill-indigo-600 text-base' : 'fill-slate-500'}`}>RB: {(pRB * 100).toFixed(1)}%</text>
                   <rect x="290" y="85" width="80" height="20" fill="transparent" />
               </g>

               <g 
                 className="cursor-pointer" 
                 onMouseEnter={() => setHoverPath('BR')} 
                 onMouseLeave={() => setHoverPath(null)}
               >
                   <text x="300" y="164" className={`text-xs font-bold transition-all ${hoverPath === 'BR' ? 'fill-indigo-600 text-base' : 'fill-slate-500'}`}>BR: {(pBR * 100).toFixed(1)}%</text>
                   <rect x="290" y="150" width="80" height="20" fill="transparent" />
               </g>

               <g 
                 className="cursor-pointer" 
                 onMouseEnter={() => setHoverPath('BB')} 
                 onMouseLeave={() => setHoverPath(null)}
               >
                   <text x="300" y="228" className={`text-xs font-bold transition-all ${hoverPath === 'BB' ? 'fill-blue-600 text-base' : 'fill-slate-500'}`}>BB: {(pBB * 100).toFixed(1)}%</text>
                   <rect x="290" y="215" width="80" height="20" fill="transparent" />
               </g>
           </svg>
       </div>

       {/* Calculation Panel */}
       <div className={`p-4 rounded-lg border text-sm mt-4 transition-colors ${hoverPath ? 'bg-amber-50 border-amber-200 text-amber-900' : 'bg-slate-50 border-slate-100 text-slate-400'}`}>
           {!hoverPath ? (
               <p className="text-center italic">Hover over an outcome (e.g., RR) to see the calculation.</p>
           ) : (
               <>
                   <p className="font-bold mb-1">
                       Calculation for <span className="font-mono bg-white px-1 rounded border border-amber-200">{hoverPath}</span> 
                       ({hoverPath[0] === 'R' ? 'Red' : 'Blue'} then {hoverPath[1] === 'R' ? 'Red' : 'Blue'}):
                   </p>
                   <div className="flex flex-wrap items-center gap-2 font-mono text-lg mt-2 justify-center sm:justify-start">
                       {/* First Draw */}
                       <span>P({hoverPath[0]})</span>
                       <span>×</span>
                       <span>P({hoverPath[1]} | {hoverPath[0]})</span>
                       <span>=</span>
                       
                       {/* Numbers */}
                       {fraction(hoverPath[0] === 'R' ? initR : initB, total)}
                       <span>×</span>
                       {fraction(
                           hoverPath === 'RR' ? r_R : hoverPath === 'RB' ? initB : hoverPath === 'BR' ? initR : b_B,
                           hoverPath[0] === 'R' ? r_Total : b_Total
                       )}
                       <span>=</span>
                       
                       {/* Result */}
                       <strong className="text-amber-700">
                           {fraction(
                               (hoverPath[0] === 'R' ? initR : initB) * (hoverPath === 'RR' ? r_R : hoverPath === 'RB' ? initB : hoverPath === 'BR' ? initR : b_B),
                               total * (hoverPath[0] === 'R' ? r_Total : b_Total)
                           )}
                       </strong>
                   </div>
                   {!replacement && hoverPath[0] === hoverPath[1] && (
                       <p className="text-xs mt-3 text-rose-600 font-bold bg-white p-2 rounded inline-block border border-rose-100">
                           ⚠ Notice: The numerator decreased because we kept the first {hoverPath[0] === 'R' ? 'Red' : 'Blue'} item!
                       </p>
                   )}
               </>
           )}
       </div>
    </div>
  );
};

export default ProbabilityTreeWidget;